CMakeLists.txt
LICENSE
MANIFEST.in
README.md
pyproject.toml
setup.py
bmtrain/__init__.py
bmtrain/block_layer.py
bmtrain/checkpointing.py
bmtrain/debug.py
bmtrain/global_var.py
bmtrain/init.py
bmtrain/layer.py
bmtrain/param_init.py
bmtrain/parameter.py
bmtrain/pipe_layer.py
bmtrain/store.py
bmtrain/synchronize.py
bmtrain/utils.py
bmtrain/wrapper.py
bmtrain.egg-info/PKG-INFO
bmtrain.egg-info/SOURCES.txt
bmtrain.egg-info/dependency_links.txt
bmtrain.egg-info/requires.txt
bmtrain.egg-info/top_level.txt
bmtrain/benchmark/__init__.py
bmtrain/benchmark/all_gather.py
bmtrain/benchmark/reduce_scatter.py
bmtrain/benchmark/send_recv.py
bmtrain/benchmark/shape.py
bmtrain/benchmark/utils.py
bmtrain/distributed/__init__.py
bmtrain/distributed/ops.py
bmtrain/inspect/__init__.py
bmtrain/inspect/format.py
bmtrain/inspect/model.py
bmtrain/inspect/tensor.py
bmtrain/loss/__init__.py
bmtrain/loss/_function.py
bmtrain/loss/cross_entropy.py
bmtrain/lr_scheduler/__init__.py
bmtrain/lr_scheduler/cosine.py
bmtrain/lr_scheduler/exponential.py
bmtrain/lr_scheduler/linear.py
bmtrain/lr_scheduler/no_decay.py
bmtrain/lr_scheduler/noam.py
bmtrain/lr_scheduler/warmup.py
bmtrain/nccl/__init__.py
bmtrain/nccl/enums.py
bmtrain/optim/__init__.py
bmtrain/optim/_function.py
bmtrain/optim/adam.py
bmtrain/optim/adam_offload.py
bmtrain/optim/optim_manager.py
cmake/FindNCCL.cmake
csrc/bind.cpp
csrc/cuda/adam_cuda.cu
csrc/cuda/cross_entropy.cu
csrc/cuda/has_inf_nan.cu
csrc/cuda/reduce.cuh
csrc/include/adam_cpu.hpp
csrc/include/bind.hpp
csrc/include/cpu_info.h
csrc/include/nccl.hpp
tests/test_all.py
tests/test_different_output_shape.py
tests/test_dropout.py
tests/test_grad_accu.py
tests/test_has_inf_nan.py
tests/test_init_parameters.py
tests/test_init_parameters_multi_gpu.py
tests/test_inspector_hidden.py
tests/test_load_ckpt.py
tests/test_loss_func.py
tests/test_middle_hidden.py
tests/test_model_wrapper.py
tests/test_multi_return.py
tests/test_nccl_backward.py
tests/test_optim.py
tests/test_optim_state.py
tests/test_other_hidden.py
tests/test_requires_grad.py
tests/test_requires_grad_multi_gpu.py
tests/test_send_recv.py
tests/test_training.py