File tree

• Src
  • .DS_Store
  • Main_Scripts
    • Chat.py
    • ColossalAI
      • .github/workflows/scripts
        • check_doc_i18n.py
        • example_checks
          • check_dispatch_inputs.py
          • check_example_weekly.py
          • detect_changed_example.py
        • generate_leaderboard_and_send_to_lark.py
        • generate_release_draft.py
        • send_message_to_lark.py
        • update_setup_for_nightly.py
      • applications
        • Colossal-LLaMA
          • colossal_llama
            • __init__.py
            • dataset
              • __init__.py
              • conversation.py
              • loader.py
              • spliced_and_tokenized_dataset.py
            • model
              • init_model.py
            • tokenizer
              • init_tokenizer.py
            • utils
              • __init__.py
              • ckpt_io.py
              • flash_attention_patch.py
              • froze.py
              • neftune_patch.py
              • stream_chat_patch.py
          • inference_example.py
          • prepare_pretrain_dataset.py
          • prepare_sft_dataset.py
          • stream_chat_example.py
          • train.py
        • ColossalChat
          • benchmarks
            • benchmark_ppo.py
            • ray
              • 1mmt_dummy.py
              • mmmt_dummy.py
          • coati
            • __init__.py
            • dataset
              • __init__.py
              • conversation.py
              • loader.py
              • tokenization_utils.py
              • utils.py
            • experience_buffer
              • __init__.py
              • base.py
              • naive.py
              • utils.py
            • experience_maker
              • __init__.py
              • base.py
              • naive.py
            • models
              • __init__.py
              • base.py
              • critic.py
              • generation.py
              • lora.py
              • loss.py
              • reward_model.py
              • utils.py
            • quant
              • __init__.py
              • llama_gptq
                • __init__.py
                • loader.py
                • model_utils.py
                • quant.py
              • utils.py
            • ray
              • __init__.py
              • callbacks
                • __init__.py
                • base.py
                • performance_evaluator.py
              • detached_replay_buffer.py
              • detached_trainer_base.py
              • detached_trainer_ppo.py
              • experience_maker_holder.py
              • lora_constructor.py
              • utils.py
            • trainer
              • __init__.py
              • base.py
              • callbacks
                • __init__.py
                • base.py
                • performance_evaluator.py
              • dpo.py
              • ppo.py
              • rm.py
              • sft.py
              • utils.py
            • utils
              • __init__.py
              • accumulative_meter.py
              • ckpt_io.py
          • examples
            • community
              • peft
                • easy_dataset.py
                • easy_models.py
                • train_peft_prompts.py
                • train_peft_sft.py
              • ray
                • ray_job_script.py
                • train_prompts_on_ray.py
            • data_preparation_scripts
              • prepare_dataset.py
            • inference
              • chatio.py
              • inference.py
              • web_chatbot
                • locustfile.py
                • server.py
                • utils.py
            • ray
              • 1mmt_prompt.py
              • mmmt_prompt.py
            • training_scripts
              • train_dpo.py
              • train_ppo.py
              • train_rm.py
              • train_sft.py
          • setup.py
          • tests
            • __init__.py
            • generate_dummy_datasets_for_testing.py
            • test_lora.py
            • verify_chat_data.py
        • ColossalEval
          • colossal_eval
            • __init__.py
            • dataset
              • __init__.py
              • agieval.py
              • base.py
              • ceval.py
              • cmmlu.py
              • colossalai.py
              • cvalues.py
              • gaokaobench.py
              • gsm.py
              • longbench.py
              • mmlu.py
              • mtbench.py
              • safetybench_en.py
              • safetybench_zh.py
            • evaluate
              • __init__.py
              • dataset_evaluator
                • __init__.py
                • dataset_evaluator.py
                • gpt_judge.py
                • metrics.py
              • evaluator.py
              • gpt_evaluate.py
              • utils.py
            • models
              • __init__.py
              • base.py
              • chatglm.py
              • huggingface.py
            • utils
              • __init__.py
              • conversation.py
              • utilities.py
          • examples
            • dataset_evaluation
              • eval_dataset.py
              • inference.py
            • gpt_evaluation
              • eval.py
              • inference.py
          • setup.py
        • ColossalMoE
          • colossal_moe
            • __init__.py
            • models
              • __init__.py
              • mixtral_checkpoint.py
              • mixtral_layer.py
              • mixtral_policy.py
            • utils.py
          • infer.py
          • setup.py
          • tests
            • __init__.py
            • test_mixtral_layer.py
            • test_moe_checkpoint.py
          • train.py
        • ColossalQA
          • colossalqa
            • __init__.py
            • chain
              • __init__.py
              • memory
                • __init__.py
                • summary.py
              • retrieval_qa
                • __init__.py
                • base.py
                • load_chain.py
                • stuff.py
            • data_loader
              • __init__.py
              • document_loader.py
              • table_dataloader.py
            • local
              • __init__.py
              • colossalcloud_llm.py
              • llm.py
              • pangu_llm.py
              • utils.py
            • memory.py
            • mylogging.py
            • prompt
              • prompt.py
            • retrieval_conversation_en.py
            • retrieval_conversation_universal.py
            • retrieval_conversation_zh.py
            • retriever.py
            • text_splitter
              • __init__.py
              • chinese_text_splitter.py
              • utils.py
            • utils.py
          • examples
            • conversation_agent_chatgpt.py
            • retrieval_conversation_chatgpt.py
            • retrieval_conversation_en.py
            • retrieval_conversation_en_customer_service.py
            • retrieval_conversation_universal.py
            • retrieval_conversation_zh.py
            • retrieval_intent_classification_zh_customer_service.py
            • webui_demo
              • RAG_ChatBot.py
              • config.py
              • server.py
              • utils.py
              • webui.py
          • setup.py
          • tests
            • __init__.py
            • test_document_loader.py
            • test_memory.py
            • test_retrieval_qa.py
            • test_text_splitter.py
      • colossalai
        • _C
          • __init__.py
        • __init__.py
        • _analyzer
          • __init__.py
          • _subclasses
            • __init__.py
            • _meta_registration.py
            • _monkey_patch.py
            • flop_tensor.py
            • meta_tensor.py
          • envs.py
          • fx
            • __init__.py
            • codegen.py
            • graph_module.py
            • node_util.py
            • passes
              • __init__.py
              • graph_profile.py
              • shape_prop.py
            • symbolic_profile.py
            • tracer
              • __init__.py
              • bias_addition.py
              • custom_leaf_module.py
              • proxy.py
              • symbolic_trace.py
              • tracer.py
        • accelerator
          • __init__.py
          • api.py
          • base_accelerator.py
          • cpu_accelerator.py
          • cuda_accelerator.py
          • npu_accelerator.py
        • amp
          • __init__.py
          • naive_amp
            • __init__.py
            • grad_scaler
              • __init__.py
              • base_grad_scaler.py
              • constant_grad_scaler.py
              • dynamic_grad_scaler.py
            • mixed_precision_mixin
              • __init__.py
              • base.py
              • bf16.py
              • fp16.py
            • mixed_precision_optimizer.py
        • auto_parallel
          • __init__.py
          • checkpoint
            • __init__.py
            • build_c_ext.py
            • ckpt_solver_base.py
            • ckpt_solver_chen.py
            • ckpt_solver_rotor.py
            • operation.py
          • meta_profiler
            • __init__.py
            • constants.py
            • meta_registry
              • __init__.py
              • activation.py
              • binary_elementwise_ops.py
              • conv.py
              • embedding.py
              • linear.py
              • non_spmd.py
              • norm.py
              • pooling.py
              • tensor.py
              • where.py
            • registry.py
            • shard_metainfo.py
          • offload
            • __init__.py
            • amp_optimizer.py
            • base_offload_module.py
            • mem_optimize.py
            • region.py
            • region_manager.py
            • runtime.py
            • solver.py
            • training_simulator.py
            • util.py
          • passes
            • __init__.py
            • comm_metainfo_pass.py
            • constants.py
            • meta_info_prop.py
            • runtime_apply_pass.py
            • runtime_preparation_pass.py
          • pipeline_shard
            • __init__.py
          • tensor_shard
            • __init__.py
            • constants.py
            • initialize.py
            • node_handler
              • __init__.py
              • addmm_handler.py
              • batch_norm_handler.py
              • binary_elementwise_handler.py
              • bmm_handler.py
              • conv_handler.py
              • default_reshape_handler.py
              • embedding_handler.py
              • getattr_handler.py
              • getitem_handler.py
              • layer_norm_handler.py
              • linear_handler.py
              • matmul_handler.py
              • node_handler.py
              • normal_pooling_handler.py
              • output_handler.py
              • permute_handler.py
              • placeholder_handler.py
              • registry.py
              • softmax_handler.py
              • split_handler.py
              • strategy
                • __init__.py
                • batch_norm_generator.py
                • binary_elementwise_generator.py
                • conv_strategy_generator.py
                • embedding_generator.py
                • getattr_generator.py
                • getitem_generator.py
                • layer_norm_generator.py
                • matmul_strategy_generator.py
                • normal_pooling_generator.py
                • output_generator.py
                • placeholder_generator.py
                • reshape_generator.py
                • softmax_generator.py
                • strategy_generator.py
                • sum_generator.py
                • tensor_constructor_generator.py
                • unary_elementwise_generator.py
                • where_generator.py
              • sum_handler.py
              • tensor_constructor_handler.py
              • transpose_handler.py
              • unary_elementwise_handler.py
              • view_handler.py
              • where_handler.py
            • options.py
            • sharding_strategy.py
            • solver
              • __init__.py
              • cost_graph.py
              • graph_analysis.py
              • solver.py
              • strategies_constructor.py
            • utils
              • __init__.py
              • broadcast.py
              • factory.py
              • misc.py
              • reshape.py
              • sharding.py
        • autochunk
          • autochunk_codegen.py
          • estimate_memory.py
          • reorder_graph.py
          • search_chunk.py
          • select_chunk.py
          • trace_flow.py
          • trace_indice.py
          • utils.py
        • booster
          • __init__.py
          • accelerator.py
          • booster.py
          • mixed_precision
            • __init__.py
            • bf16.py
            • fp16_apex.py
            • fp16_naive.py
            • fp16_torch.py
            • fp8.py
            • mixed_precision_base.py
          • plugin
            • __init__.py
            • dp_plugin_base.py
            • gemini_plugin.py
            • hybrid_parallel_plugin.py
            • low_level_zero_plugin.py
            • moe_hybrid_parallel_plugin.py
            • plugin_base.py
            • pp_plugin_base.py
            • torch_ddp_plugin.py
            • torch_fsdp_plugin.py
        • checkpoint_io
          • __init__.py
          • checkpoint_io_base.py
          • general_checkpoint_io.py
          • hybrid_parallel_checkpoint_io.py
          • index_file.py
          • utils.py
        • cli
          • __init__.py
          • check
            • __init__.py
            • check_installation.py
          • cli.py
          • launcher
            • __init__.py
            • hostinfo.py
            • multinode_runner.py
            • run.py
        • cluster
          • __init__.py
          • device_mesh_manager.py
          • dist_coordinator.py
          • process_group_manager.py
          • process_group_mesh.py
        • context
          • __init__.py
          • config.py
          • singleton_meta.py
        • device
          • __init__.py
          • alpha_beta_profiler.py
          • calc_pipeline_strategy.py
          • device_mesh.py
        • fx
          • __init__.py
          • _compatibility.py
          • _meta_regist_12.py
          • _meta_regist_13.py
          • codegen
            • __init__.py
            • activation_checkpoint_codegen.py
          • graph_module.py
          • passes
            • __init__.py
            • adding_split_node_pass.py
            • concrete_info_prop.py
            • experimental
              • adding_shape_consistency_pass.py
            • meta_info_prop.py
            • passes_for_gpt2_test.py
            • shard_1d_pass.py
            • split_module.py
            • utils.py
          • profiler
            • __init__.py
            • constants.py
            • dataflow.py
            • experimental
              • __init__.py
              • constants.py
              • profiler.py
              • profiler_function
                • __init__.py
                • activation_function.py
                • arithmetic.py
                • embedding.py
                • linear.py
                • normalization.py
                • pooling.py
                • python_ops.py
                • torch_ops.py
              • profiler_module
                • __init__.py
                • activation_function.py
                • attention.py
                • convolution.py
                • dropout.py
                • embedding.py
                • linear.py
                • normalization.py
                • pooling.py
                • rnn.py
                • torch_op.py
              • registry.py
              • shard_utils.py
            • memory_utils.py
            • opcount.py
            • profiler.py
            • shard_utils.py
            • tensor.py
          • proxy.py
          • tracer
            • __init__.py
            • _meta_trace.py
            • _symbolic_trace.py
            • _tracer_utils.py
            • bias_addition_patch
              • __init__.py
              • patched_bias_addition_function
                • __init__.py
                • addbmm.py
                • addmm.py
                • bias_addition_function.py
                • linear.py
              • patched_bias_addition_module
                • __init__.py
                • bias_addition_module.py
                • conv.py
                • linear.py
            • experimental.py
            • meta_patch
              • __init__.py
              • patched_function
                • __init__.py
                • activation_function.py
                • arithmetic.py
                • convolution.py
                • embedding.py
                • normalization.py
                • python_ops.py
                • torch_ops.py
              • patched_module
                • __init__.py
                • activation_function.py
                • convolution.py
                • embedding.py
                • linear.py
                • normalization.py
                • pooling.py
                • rnn.py
            • registry.py
            • tracer.py
        • inference
          • __init__.py
          • engine
            • __init__.py
            • engine.py
            • microbatch_manager.py
            • modeling
              • __init__.py
              • _utils.py
              • bloom.py
              • chatglm2.py
              • llama.py
            • policies
              • __init__.py
              • bloom.py
              • chatglm2.py
              • llama.py
          • kv_cache
            • __init__.py
            • batch_infer_state.py
            • kvcache_manager.py
          • quant
            • __init__.py
            • gptq
              • __init__.py
              • cai_gptq
                • __init__.py
                • cai_quant_linear.py
                • gptq_op.py
              • gptq_manager.py
            • smoothquant
              • __init__.py
              • models
                • __init__.py
                • base_model.py
                • linear.py
                • llama.py
                • parallel_linear.py
        • initialize.py
        • interface
          • __init__.py
          • model.py
          • optimizer.py
          • pretrained.py
        • kernel
          • __init__.py
          • jit
            • __init__.py
            • bias_dropout_add.py
            • bias_gelu.py
            • option.py
          • kernel_loader.py
          • triton
            • __init__.py
            • context_attention.py
            • copy_kv_cache_dest.py
            • custom_autotune.py
            • flash_decoding.py
            • fused_layernorm.py
            • gptq_triton.py
            • int8_rotary_embedding_kernel.py
            • llama_act_combine_kernel.py
            • qkv_matmul_kernel.py
            • self_attention_nofusion.py
            • smooth_attention.py
            • softmax.py
            • token_attention_kernel.py
        • lazy
          • __init__.py
          • construction.py
          • lazy_init.py
          • pretrained.py
        • legacy
          • __init__.py
          • amp
            • __init__.py
            • amp_type.py
            • apex_amp
              • __init__.py
              • apex_amp.py
            • naive_amp
              • __init__.py
              • _fp16_optimizer.py
              • _utils.py
              • naive_amp.py
            • torch_amp
              • __init__.py
              • _grad_scaler.py
              • torch_amp.py
          • builder
            • __init__.py
            • builder.py
          • communication
            • __init__.py
            • collective.py
            • p2p.py
            • p2p_v2.py
            • ring.py
            • utils.py
          • constants.py
          • context
            • __init__.py
            • parallel_context.py
            • parallel_mode.py
            • process_group_initializer
              • __init__.py
              • initializer_1d.py
              • initializer_2d.py
              • initializer_2p5d.py
              • initializer_3d.py
              • initializer_data.py
              • initializer_model.py
              • initializer_pipeline.py
              • initializer_sequence.py
              • initializer_tensor.py
              • process_group_initializer.py
            • random
              • __init__.py
              • _helper.py
              • seed_manager.py
          • core.py
          • engine
            • __init__.py
            • _base_engine.py
            • gradient_accumulation
              • __init__.py
              • _gradient_accumulation.py
            • gradient_handler
              • __init__.py
              • _base_gradient_handler.py
              • _data_parallel_gradient_handler.py
              • _moe_gradient_handler.py
              • _pipeline_parallel_gradient_handler.py
              • _sequence_parallel_gradient_handler.py
              • _zero_gradient_handler.py
              • utils.py
            • schedule
              • __init__.py
              • _base_schedule.py
              • _non_pipeline_schedule.py
              • _pipeline_schedule.py
              • _pipeline_schedule_v2.py
          • global_variables.py
          • inference
            • __init__.py
            • async_engine.py
            • async_manager.py
            • dynamic_batching
              • __init__.py
              • get_tokenizer.py
              • infer_batch.py
              • io_struct.py
              • ray_dist_init.py
              • ray_init_config.py
              • req_queue.py
              • sampling_params.py
              • stats.py
            • hybridengine
              • __init__.py
              • engine.py
              • modeling
                • __init__.py
                • _utils.py
                • llama.py
              • polices
                • __init__.py
                • llama.py
            • manager.py
            • pipeline
              • __init__.py
              • benchmark
                • benchmark.py
              • microbatch_manager.py
            • quant
              • gptq
                • __init__.py
                • cai_gptq
                  • __init__.py
                  • cai_quant_linear.py
                  • gptq_op.py
              • smoothquant
                • __init__.py
                • models
                  • __init__.py
                  • base_model.py
                  • linear.py
                  • llama.py
            • serving
              • ray_serve
                • Colossal_Inference_rayserve.py
                • send_request.py
                • send_requests.py
              • torch_serve
                • Colossal_Inference_Handler.py
            • tensor_parallel
              • __init__.py
              • batch_infer_state.py
              • engine.py
              • kvcache_manager.py
              • modeling
                • __init__.py
                • _utils.py
                • bloom.py
                • chatglm2.py
                • llama.py
              • policies
                • __init__.py
                • bloom.py
                • chatglm2.py
                • llama.py
          • initialize.py
          • nn
            • __init__.py
            • _ops
              • __init__.py
              • _utils.py
            • layer
              • __init__.py
              • base_layer.py
              • colossalai_layer
                • __init__.py
                • _utils.py
                • dropout.py
                • embedding.py
                • linear.py
                • normalization.py
              • parallel_1d
                • __init__.py
                • _operation.py
                • _utils.py
                • layers.py
              • parallel_2d
                • __init__.py
                • _operation.py
                • _utils.py
                • layers.py
              • parallel_2p5d
                • __init__.py
                • _operation.py
                • _utils.py
                • layers.py
              • parallel_3d
                • __init__.py
                • _operation.py
                • _utils.py
                • layers.py
              • parallel_sequence
                • __init__.py
                • _operation.py
                • _utils.py
                • layers.py
              • utils
                • __init__.py
                • common.py
              • vanilla
                • __init__.py
                • layers.py
              • wrapper
                • __init__.py
                • pipeline_wrapper.py
            • loss
              • __init__.py
              • loss_1d.py
              • loss_2d.py
              • loss_2p5d.py
              • loss_3d.py
            • metric
              • __init__.py
              • _utils.py
              • accuracy_2d.py
              • accuracy_2p5d.py
              • accuracy_3d.py
            • parallel
              • __init__.py
              • data_parallel.py
              • layers
                • __init__.py
                • cache_embedding
                  • __init__.py
                  • base_embedding.py
                  • cache_mgr.py
                  • cached_embedding.py
                  • copyer.py
                  • embedding_config.py
                  • parallel_cached_embedding.py
                  • parallel_cached_embedding_tablewise.py
                  • parallel_cached_embedding_tablewise_split_cache.py
                • colo_module.py
                • embedding.py
                • linear.py
                • module_utils.py
              • reducer.py
          • pipeline
            • __init__.py
            • layer_spec.py
            • middleware
              • __init__.py
              • adaptor
                • __init__.py
                • fx.py
              • topo.py
            • pipelinable.py
            • pipeline_process_group.py
            • rpc
              • __init__.py
              • _pipeline_base.py
              • _pipeline_schedule.py
              • utils.py
            • utils.py
          • registry
            • __init__.py
            • registry.py
          • tensor
            • __init__.py
            • compute_spec.py
            • const.py
            • dist_spec_mgr.py
            • distspec.py
            • op_wrapper.py
            • process_group.py
            • tensor_spec.py
          • trainer
            • __init__.py
            • _trainer.py
            • hooks
              • __init__.py
              • _base_hook.py
              • _checkpoint_hook.py
              • _commons_.py
              • _log_hook.py
              • _lr_scheduler_hook.py
              • _metric_hook.py
          • utils
            • __init__.py
            • activation_checkpoint.py
            • checkpointing.py
            • checkpoint
              • __init__.py
              • module_checkpoint.py
              • utils.py
            • common.py
            • data_sampler
              • __init__.py
              • base_sampler.py
              • data_parallel_sampler.py
            • memory.py
            • profiler
              • __init__.py
              • extention.py
              • legacy
                • __init__.py
                • comm_profiler.py
                • pcie_profiler.py
                • prof_utils.py
              • profiler.py
              • stateful_tensor_mem_extention.py
          • zero
            • __init__.py
            • gemini
              • __init__.py
              • colo_init_context.py
              • gemini_context.py
              • ophooks
                • __init__.py
                • _shard_grad_ophook.py
                • _shard_param_ophook.py
                • runtime_mem_tracer_hook.py
                • utils.py
              • paramhooks
                • __init__.py
                • _param_hookmgr.py
              • stateful_tensor.py
              • stateful_tensor_mgr.py
              • tensor_placement_policy.py
              • tensor_utils.py
            • init_ctx
              • __init__.py
              • init_context.py
            • shard_utils
              • __init__.py
              • base_shard_strategy.py
              • bucket_tensor_shard_strategy.py
              • commons.py
              • tensor_shard_strategy.py
            • sharded_model
              • __init__.py
              • _utils.py
              • reduce_scatter.py
              • sharded_model_v2.py
              • utils.py
              • zero_hook.py
            • sharded_optim
              • __init__.py
              • sharded_optim_v2.py
            • sharded_param
              • __init__.py
              • sharded_param.py
              • sharded_tensor.py
        • logging
          • __init__.py
          • logger.py
        • moe
          • __init__.py
          • _operation.py
          • checkpoint.py
          • experts.py
          • layers.py
          • load_balance.py
          • loss.py
          • manager.py
          • routers.py
          • utils.py
        • nn
          • __init__.py
          • init.py
          • layer
            • __init__.py
            • layernorm.py
            • scaled_softmax.py
            • utils.py
          • loss
            • __init__.py
          • lr_scheduler
            • __init__.py
            • cosine.py
            • delayed.py
            • linear.py
            • multistep.py
            • onecycle.py
            • poly.py
            • torch.py
          • optimizer
            • __init__.py
            • cpu_adam.py
            • fused_adam.py
            • fused_lamb.py
            • fused_sgd.py
            • hybrid_adam.py
            • lamb.py
            • lars.py
            • nvme_optimizer.py
        • pipeline
          • __init__.py
          • p2p.py
          • schedule
            • __init__.py
            • _utils.py
            • base.py
            • generate.py
            • interleaved_pp.py
            • one_f_one_b.py
          • stage_manager.py
        • shardformer
          • __init__.py
          • _utils.py
          • examples
            • convergence_benchmark.py
            • data.py
            • performance_benchmark.py
          • layer
            • __init__.py
            • _operation.py
            • attn.py
            • dropout.py
            • embedding.py
            • linear.py
            • loss.py
            • normalization.py
            • parallel_module.py
            • qkv_fused_linear.py
            • utils.py
          • modeling
            • __init__.py
            • bert.py
            • blip2.py
            • bloom.py
            • chatglm2.py
            • chatglm2_6b
              • __init__.py
              • configuration_chatglm.py
              • modeling_chatglm.py
            • falcon.py
            • gpt2.py
            • gptj.py
            • jit.py
            • llama.py
            • mistral.py
            • opt.py
            • sam.py
            • t5.py
            • vit.py
            • whisper.py
          • policies
            • __init__.py
            • auto_policy.py
            • base_policy.py
            • bert.py
            • blip2.py
            • bloom.py
            • chatglm2.py
            • falcon.py
            • gpt2.py
            • gptj.py
            • llama.py
            • mistral.py
            • opt.py
            • sam.py
            • t5.py
            • vit.py
            • whisper.py
          • shard
            • __init__.py
            • grad_ckpt_config.py
            • shard_config.py
            • sharder.py
            • shardformer.py
            • utils.py
        • tensor
          • __init__.py
          • colo_parameter.py
          • colo_tensor.py
          • comm_spec.py
          • d_tensor
            • __init__.py
            • api.py
            • comm_spec.py
            • layout.py
            • layout_converter.py
            • misc.py
            • sharding_spec.py
            • utils.py
          • moe_tensor
            • __init__.py
            • api.py
            • moe_info.py
          • padded_tensor
            • __init__.py
            • api.py
          • param_op_hook.py
          • shape_consistency.py
          • sharding_spec.py
          • utils.py
        • testing
          • __init__.py
          • comparison.py
          • pytest_wrapper.py
          • random.py
          • utils.py
        • utils
          • __init__.py
          • common.py
          • memory.py
          • model
            • __init__.py
            • utils.py
          • multi_tensor_apply
            • __init__.py
            • multi_tensor_apply.py
          • rank_recorder
            • __init__.py
            • rank_recorder.py
          • tensor_detector
            • __init__.py
            • tensor_detector.py
          • timer.py
        • zero
          • __init__.py
          • gemini
            • __init__.py
            • chunk
              • __init__.py
              • chunk.py
              • manager.py
              • search_utils.py
              • utils.py
            • gemini_ddp.py
            • gemini_hook.py
            • gemini_mgr.py
            • gemini_optimizer.py
            • memory_tracer
              • __init__.py
              • chunk_memstats_collector.py
              • memory_monitor.py
              • memory_stats.py
              • memstats_collector.py
              • param_runtime_order.py
              • runtime_mem_tracer.py
              • static_memstats_collector.py
              • utils.py
            • placement_policy.py
            • utils.py
          • low_level
            • __init__.py
            • _utils.py
            • bookkeeping
              • __init__.py
              • base_store.py
              • bucket_store.py
              • gradient_store.py
              • parameter_store.py
              • tensor_bucket.py
            • low_level_optim.py
          • wrapper.py
      • examples
        • __init__.py
        • community
          • fp8/mnist
            • main.py
          • roberta
            • preprocessing
              • get_mask.py
              • sentence_split.py
              • tokenize_mask.py
            • pretraining
              • arguments.py
              • bert_dataset_provider.py
              • evaluation.py
              • loss.py
              • model
                • bert.py
                • deberta_v2.py
              • nvidia_bert_dataset_provider.py
              • pretrain_utils.py
              • run_pretraining.py
              • utils
                • WandbLog.py
                • exp_util.py
                • global_vars.py
                • logger.py
        • images
          • diffusion
            • ldm
              • data
                • __init__.py
                • base.py
                • cifar10.py
                • imagenet.py
                • lsun.py
                • teyvat.py
              • lr_scheduler.py
              • models
                • autoencoder.py
                • diffusion
                  • __init__.py
                  • classifier.py
                  • ddim.py
                  • ddpm.py
                  • dpm_solver
                    • __init__.py
                    • dpm_solver.py
                    • sampler.py
                  • plms.py
                  • sampling_util.py
              • modules
                • attention.py
                • diffusionmodules
                  • __init__.py
                  • model.py
                  • openaimodel.py
                  • upscaling.py
                  • util.py
                • distributions
                  • __init__.py
                  • distributions.py
                • ema.py
                • encoders
                  • __init__.py
                  • modules.py
                • image_degradation
                  • __init__.py
                  • bsrgan.py
                  • bsrgan_light.py
                  • utils_image.py
                • midas
                  • __init__.py
                  • api.py
                  • midas
                    • __init__.py
                    • base_model.py
                    • blocks.py
                    • dpt_depth.py
                    • midas_net.py
                    • midas_net_custom.py
                    • transforms.py
                    • vit.py
                  • utils.py
              • util.py
            • main.py
            • scripts
              • img2img.py
              • inpaint.py
              • knn2img.py
              • sample_diffusion.py
              • tests
                • test_checkpoint.py
                • test_watermark.py
              • train_searcher.py
              • txt2img.py
              • utils.py
            • setup.py
          • dreambooth
            • debug.py
            • inference.py
            • train_dreambooth.py
            • train_dreambooth_colossalai.py
            • train_dreambooth_colossalai_lora.py
            • train_dreambooth_inpaint.py
          • resnet
            • eval.py
            • train.py
          • vit
            • args.py
            • data.py
            • vit_benchmark.py
            • vit_train_demo.py
        • inference
          • benchmark_llama.py
          • build_smoothquant_weight.py
          • run_llama_inference.py
        • language
          • __init__.py
          • bert
            • benchmark.py
            • benchmark_utils.py
            • data.py
            • finetune.py
          • commons
            • utils.py
          • data_utils.py
          • gpt
            • experiments
              • auto_offload
                • model_zoo.py
                • train_gpt_offload.py
              • auto_parallel
                • auto_parallel_with_gpt.py
                • gpt_modules.py
              • pipeline_parallel
                • model_zoo.py
                • train_gpt_pp.py
            • gemini
              • commons
                • model_zoo.py
                • utils.py
              • train_gpt_demo.py
            • hybridparallelism
              • benchmark.py
              • data.py
              • finetune.py
            • titans
              • configs
                • gpt2_small_zero3_pp1d.py
                • gpt3_zero3_pp1d.py
              • dataset
                • webtext.py
              • model
                • __init__.py
                • embed.py
                • gpt1d.py
                • pipeline_gpt1d.py
              • train_gpt.py
          • grok-1
            • grok1_policy.py
            • inference.py
            • inference_tp.py
            • utils.py
          • llama
            • benchmark.py
          • model_utils.py
          • openmoe
            • benchmark
              • benchmark_cai.py
              • benchmark_fsdp.py
              • utils.py
            • infer.py
            • model
              • __init__.py
              • convert_openmoe_ckpt.py
              • modeling_openmoe.py
              • openmoe_policy.py
            • train.py
          • opt
            • args.py
            • data.py
            • opt_benchmark.py
            • opt_train_demo.py
          • palm
            • palm_pytorch
              • __init__.py
              • autoregressive_wrapper.py
              • palm_pytorch.py
            • train.py
          • performance_evaluator.py
        • tutorial
          • auto_parallel
            • auto_ckpt_batchsize_test.py
            • auto_ckpt_solver_test.py
            • auto_parallel_with_resnet.py
            • bench_utils.py
            • config.py
            • setup.py
          • download_cifar10.py
          • hybrid_parallel
            • config.py
            • train.py
          • large_batch_optimizer
            • config.py
            • train.py
          • new_api
            • cifar_resnet
              • eval.py
              • train.py
            • cifar_vit
              • train.py
            • glue_bert
              • data.py
              • finetune.py
          • opt
            • inference
              • batch.py
              • benchmark
                • locustfile.py
              • cache.py
              • opt_fastapi.py
              • opt_server.py
              • script
                • process-opt-175b
                  • convert_ckpt.py
                • processing_ckpt_66b.py
            • opt
              • colossalai_zero.py
              • context.py
              • run_clm.py
          • sequence_parallel
            • config.py
            • data
              • __init__.py
              • bert_helper.py
              • datasets
                • __init__.py
                • bert_dataset.py
                • blendable_dataset.py
                • builder.py
                • data_samplers.py
                • dataset_utils.py
                • ict_dataset.py
                • indexed_dataset.py
                • test
                  • test_indexed_dataset.py
              • dummy_dataloader.py
              • tokenizer
                • __init__.py
                • bert_tokenization.py
                • tokenizer.py
            • loss_func
              • __init__.py
              • bert_loss.py
              • cross_entropy.py
              • utils.py
            • lr_scheduler
              • __init__.py
              • annealing_lr.py
            • model
              • __init__.py
              • bert.py
              • layers
                • __init__.py
                • bert_layer.py
                • dropout.py
                • embedding.py
                • head.py
                • init_method.py
                • linear.py
                • mlp.py
                • pooler.py
                • preprocess.py
            • train.py
      • extensions
        • __init__.py
        • base_extension.py
        • cpp_extension.py
        • cpu_adam
          • __init__.py
          • cpu_adam_arm.py
          • cpu_adam_x86.py
        • csrc
          • __init__.py
          • cuda
            • layer_norm_cuda_kernel.cu
            • moe_cuda_kernel.cu
            • multi_tensor_adam.cu
            • multi_tensor_l2norm_kernel.cu
            • multi_tensor_lamb.cu
            • multi_tensor_scale_kernel.cu
            • multi_tensor_sgd_kernel.cu
            • scaled_masked_softmax_cuda.cu
            • scaled_upper_triang_masked_softmax_cuda.cu
          • scaled_softmax.py
        • cuda_extension.py
        • flash_attention
          • __init__.py
          • flash_attention_dao_cuda.py
          • flash_attention_npu.py
          • flash_attention_sdpa_cuda.py
        • layernorm
          • __init__.py
          • layernorm_cuda.py
        • moe
          • __init__.py
          • moe_cuda.py
        • optimizer
          • __init__.py
          • fused_optimizer_cuda.py
        • softmax
          • __init__.py
          • scaled_masked_softmax_cuda.py
          • scaled_upper_triangle_masked_softmax_cuda.py
        • triton_extension.py
        • utils.py
      • setup.py
      • tests
        • __init__.py
        • kit
          • __init__.py
          • model_zoo
            • __init__.py
            • custom
              • __init__.py
              • base.py
              • hanging_param_model.py
              • nested_model.py
              • repeated_computed_layers.py
              • simple_net.py
            • diffusers
              • __init__.py
              • diffusers.py
            • executor.py
            • registry.py
            • timm
              • __init__.py
              • timm.py
            • torchaudio
              • __init__.py
              • torchaudio.py
            • torchrec
              • __init__.py
              • torchrec.py
            • torchvision
              • __init__.py
              • torchvision.py
            • transformers
              • __init__.py
              • albert.py
              • bert.py
              • blip2.py
              • bloom.py
              • chatglm2.py
              • falcon.py
              • gpt.py
              • gptj.py
              • llama.py
              • mistral.py
              • opt.py
              • sam.py
              • t5.py
              • vit.py
              • whisper.py
        • test_analyzer
          • __init__.py
          • test_fx
            • __init__.py
            • test_bias_addition.py
            • test_mod_dir.py
            • test_nested_ckpt.py
            • test_shape_prop.py
            • test_symbolic_profile.py
            • zoo.py
          • test_subclasses
            • __init__.py
            • test_aten.py
            • test_flop_tensor.py
            • test_meta_mode.py
        • test_auto_parallel
          • __init__.py
          • test_ckpt_solvers
            • test_C_solver_consistency.py
            • test_ckpt_torchvision.py
            • test_linearize.py
          • test_offload
            • model_utils.py
            • test_perf.py
            • test_solver.py
          • test_pass
            • __init__.py
            • test_node_converting_pass.py
            • test_size_value_converting_pass.py
          • test_tensor_shard
            • __init__.py
            • test_bias_addition_forward.py
            • test_broadcast.py
            • test_checkpoint.py
            • test_compatibility_with_ddp.py
            • test_compatibility_with_gemini.py
            • test_find_repeat_block.py
            • test_gpt
              • __init__.py
              • gpt_modules.py
              • test_runtime_with_gpt_modules.py
              • test_solver_with_gpt_module.py
            • test_liveness_analysis.py
            • test_metainfo
              • test_activation_metainfo.py
              • test_binary_elementwise_metainfo.py
              • test_conv_metainfo.py
              • test_embedding_metainfo.py
              • test_linear_metainfo.py
              • test_matmul_metainfo.py
              • test_norm_metainfo.py
              • test_pooling_metainfo.py
              • test_tensor_metainfo.py
              • test_where_metainfo.py
              • utils.py
            • test_node_handler
              • __init__.py
              • test_addbmm_handler.py
              • test_addmm_handler.py
              • test_batch_norm_handler.py
              • test_bias_linear_function_node.py
              • test_bias_linear_module_node.py
              • test_binary_elementwise_handler.py
              • test_bmm_handler.py
              • test_conv_handler.py
              • test_default_reshape_handler.py
              • test_embedding_handler.py
              • test_getattr_handler.py
              • test_getitem_handler.py
              • test_layer_norm_handler.py
              • test_linear_handler.py
              • test_matmul_handler.py
              • test_norm_pooling_handler.py
              • test_output_handler.py
              • test_permute_and_transpose_handler.py
              • test_placeholder_handler.py
              • test_shard_option.py
              • test_softmax_handler.py
              • test_split_handler.py
              • test_sum_handler.py
              • test_tensor_constructor.py
              • test_unary_element_wise_handler.py
              • test_view_handler.py
              • test_where_handler.py
              • utils.py
            • test_solver_with_resnet_v2.py
        • test_autochunk
          • test_autochunk_alphafold
            • benchmark_autochunk_alphafold.py
            • test_autochunk_alphafold_utils.py
            • test_autochunk_evoformer_block.py
            • test_autochunk_evoformer_stack.py
            • test_autochunk_extramsa_block.py
          • test_autochunk_diffuser
            • benchmark_autochunk_diffuser.py
            • test_autochunk_diffuser_utils.py
            • test_autochunk_unet.py
          • test_autochunk_transformer
            • benchmark_autochunk_transformer.py
            • test_autochunk_gpt.py
            • test_autochunk_transformer_utils.py
          • test_autochunk_vit
            • test_autochunk_vit.py
            • test_autochunk_vit_utils.py
        • test_booster
          • test_accelerator.py
          • test_mixed_precision
            • test_fp16_torch.py
          • test_plugin
            • test_3d_plugin.py
            • test_dp_plugin_base.py
            • test_gemini_plugin.py
            • test_low_level_zero_plugin.py
            • test_torch_ddp_plugin.py
            • test_torch_fsdp_plugin.py
        • test_checkpoint_io
          • test_gemini_checkpoint_io.py
          • test_gemini_torch_compability.py
          • test_general_checkpoint_io.py
          • test_hybrid_parallel_plugin_checkpoint_io.py
          • test_low_level_zero_checkpoint_io.py
          • test_plugins_huggingface_compatibility.py
          • test_torch_ddp_checkpoint_io.py
          • test_torch_fsdp_checkpoint_io.py
          • utils.py
        • test_cluster
          • test_device_mesh_manager.py
          • test_process_group_mesh.py
        • test_config
          • sample_config.py
          • test_load_config.py
        • test_device
          • test_alpha_beta.py
          • test_device_mesh.py
          • test_extract_alpha_beta.py
          • test_init_logical_pg.py
          • test_search_logical_device_mesh.py
        • test_fx
          • test_codegen
            • test_activation_checkpoint_codegen.py
            • test_nested_activation_checkpoint_codegen.py
            • test_offload_codegen.py
          • test_coloproxy.py
          • test_comm_size_compute.py
          • test_graph_manipulation.py
          • test_meta_info_prop.py
          • test_meta
            • test_aten.py
            • test_backward.py
            • test_meta_trace.py
          • test_parallel_1d.py
          • test_pipeline_passes.py
          • test_pipeline
            • test_hf_model
              • hf_utils.py
              • test_albert.py
              • test_bert.py
              • test_gpt.py
              • test_opt.py
              • test_t5.py
            • test_timm_model
              • test_timm.py
              • timm_utils.py
            • test_topo
              • test_topo.py
              • topo_utils.py
            • test_torchvision
              • test_torchvision.py
          • test_profiler
            • gpt_utils.py
            • test_profiler_meta_info_prop.py
          • test_tracer
            • test_activation_checkpoint_annotation.py
            • test_bias_addition_module.py
            • test_control_flow.py
            • test_functional_conv.py
            • test_hf_model
              • hf_tracer_utils.py
              • test_hf_albert.py
              • test_hf_bert.py
              • test_hf_diffuser.py
              • test_hf_gpt.py
              • test_hf_opt.py
              • test_hf_t5.py
            • test_patched_module.py
            • test_patched_op.py
            • test_timm_model
              • test_timm_model.py
            • test_torchaudio_model
              • test_torchaudio_model.py
              • torchaudio_utils.py
            • test_torchrec_model
              • test_deepfm_model.py
              • test_dlrm_model.py
            • test_torchvision_model
              • test_torchvision_model.py
        • test_gptq
          • test_gptq_linear.py
        • test_infer
          • _utils.py
          • test_hybrid_bloom.py
          • test_hybrid_chatglm2.py
          • test_hybrid_llama.py
          • test_kvcache_manager.py
        • test_lazy
          • lazy_init_utils.py
          • test_from_pretrained.py
          • test_models.py
          • test_ops.py
        • test_legacy
          • test_amp
            • test_naive_fp16.py
            • test_torch_fp16.py
          • test_comm
            • test_boardcast_send_recv_v2.py
            • test_comm.py
            • test_object_list_p2p.py
            • test_object_list_p2p_v2.py
          • test_context
            • configs
              • parallel_2d_init.py
              • parallel_2p5d_init.py
              • parallel_3d_init.py
            • test_hybrid_parallel.py
          • test_data
            • test_cifar10_dataset.py
            • test_data_parallel_sampler.py
            • test_deterministic_dataloader.py
          • test_engine
            • test_engine.py
            • test_gradient_accumluation.py
          • test_layers
            • test_1d
              • checks_1d
                • __init__.py
                • check_layer_1d.py
                • common.py
              • test_1d.py
            • test_2d
              • checks_2d
                • __init__.py
                • check_layer_2d.py
                • check_operation_2d.py
                • common.py
              • test_2d.py
            • test_2p5d
              • checks_2p5d
                • __init__.py
                • check_layer_2p5d.py
                • check_operation_2p5d.py
                • common.py
              • test_2p5d.py
            • test_3d
              • checks_3d
                • __init__.py
                • check_layer_3d.py
                • common.py
              • test_3d.py
            • test_cache_embedding.py
            • test_sequence
              • checks_seq
                • __init__.py
                • check_layer_seq.py
              • test_sequence.py
          • test_pipeline
            • rpc_test_utils.py
            • test_cuda_rpc_chimera.py
            • test_cuda_rpc_optimizer.py
            • test_cuda_rpc_pipeline.py
            • test_cuda_rpc_value_correctness.py
            • test_middleware_1f1b.py
            • test_pipelinable.py
            • test_pipeline_process_group.py
          • test_tensor
            • common_utils
              • __init__.py
              • _utils.py
            • core
              • test_dist_spec_mgr.py
            • test_parameter.py
          • test_trainer
            • test_pipeline
              • test_p2p.py
              • test_pipeline_schedule.py
            • test_trainer_with_non_pipe_schedule.py
            • test_trainer_with_pipe_schedule.py
          • test_utils
            • test_activation_checkpointing.py
            • test_checkpoint
              • test_checkpoint_1d.py
              • test_checkpoint_2d.py
              • test_checkpoint_2p5d.py
              • test_checkpoint_3d.py
            • test_memory.py
            • test_norm_gradient_clipping.py
          • test_zero
            • test_commons.py
        • test_moe
          • moe_utils.py
          • test_grad_handler.py
          • test_kernel.py
          • test_moe_checkpoint.py
          • test_moe_ep_tp.py
          • test_moe_group.py
          • test_moe_hybrid_zero.py
          • test_moe_load_balance.py
          • test_moe_router.py
          • test_moe_zero_fwd_bwd.py
          • test_moe_zero_optim.py
        • test_optimizer
          • test_adam_kernel.py
          • test_adam_optim.py
          • test_lr_scheduler.py
          • test_nvme.py
        • test_pipeline
          • test_p2p_communication.py
          • test_pipeline_utils
            • test_t5_pipeline_utils.py
            • test_whisper_pipeline_utils.py
          • test_schedule
            • test_interleaved.py
            • test_oneF_oneB.py
            • test_pipeline_schedule_utils.py
          • test_stage_manager.py
        • test_shardformer
          • __init__.py
          • test_flash_attention.py
          • test_hybrid_parallel_grad_clip_norm
            • test_amp_optimizer.py
            • test_naive_optimizer.py
            • test_zero_optimizer.py
          • test_layer
            • test_dist_crossentropy.py
            • test_dropout.py
            • test_embedding.py
            • test_gpt2_qkv_fused_linear_1d.py
            • test_layernorm.py
            • test_linear_1d.py
            • test_qkv_fused_linear_1d.py
            • test_sequence_parallel.py
            • test_vocab_parallel_embedding_1d.py
          • test_model
            • __init__.py
            • _utils.py
            • test_shard_bert.py
            • test_shard_blip2.py
            • test_shard_bloom.py
            • test_shard_chatglm2.py
            • test_shard_falcon.py
            • test_shard_gpt2.py
            • test_shard_gptj.py
            • test_shard_llama.py
            • test_shard_mistral.py
            • test_shard_opt.py
            • test_shard_sam.py
            • test_shard_t5.py
            • test_shard_vit.py
            • test_shard_whisper.py
          • test_shard_utils.py
          • test_with_torch_ddp.py
        • test_smoothquant
          • test_llama_attention.py
          • test_llama_mlp.py
          • test_smoothquant_linear.py
          • test_sq_rotary_embedding.py
        • test_tensor
          • test_comm_spec_apply.py
          • test_dtensor
            • test_comm_spec.py
            • test_dtensor.py
            • test_dtensor_sharding_spec.py
            • test_layout_converter.py
          • test_mix_gather.py
          • test_padded_tensor.py
          • test_shape_consistency.py
          • test_shape_consistency_apply.py
          • test_sharding_spec.py
        • test_zero
          • test_gemini
            • test_chunk_mgrv2.py
            • test_chunkv2.py
            • test_fwd_bwd.py
            • test_gemini_use_rmt.py
            • test_grad_accum.py
            • test_grad_clip.py
            • test_inference.py
            • test_optim.py
            • test_runtime_mem_tracer.py
            • test_search.py
            • test_zeroddp_state_dict.py
            • test_zerooptim_state_dict.py
          • test_low_level
            • test_grad_acc.py
            • test_zero1_2.py
            • test_zero_ckpt.py
    • Dataset_download.py
    • Main.py
    • backend
      • backend_colossalai.py
      • backend_deepspeed.py
      • backend_fsdp.py
    • config
      • __init__.py
      • config_manager.py
    • core
      • __init__.py
      • benchmark.py
      • check_cuda_ops.py
      • cuda_opt_wrapper.py
      • dataset.py
      • diagnose.py
      • metal_opt_wrapper.py
      • model.py
      • moe_cuda_ops.cu
      • moe_cuda_wrapper.py
      • moe_metal_wrapper.py
      • setup.py
      • tokenizer.py
      • transformer_ops.cu
      • unified_ops.py
    • cuda_debug_script.py
    • debug_overfit.py
    • deepspeed_integration.py
    • moe_inference_runtime
      • backends
        • cuda_kernels.cu
      • chat.py
      • check_config.py
      • decode_tokens.py
      • export
        • export_model.py
    • monitoring
      • __init__.py
      • logger.py
    • multi_source_dataset.py
    • router_training_system.py
    • security
      • auth.py
      • input_validator.py
      • rate_limiter.py
    • test_data.py
    • test_mps.py
    • training
      • __init__.py
      • benchmark_cuda_kernels.py
      • checkpoint.py
      • chinchilla_scaler.py
      • cuda_kernels.py
      • fused_grad_clip.cu
      • fused_loss.cu
      • orchestrator.py
      • test_grad_clip_colab.py
      • trainer.py
      • training_loop.py
      • verify_grad_clip_fix.py
    • utils
      • __init__.py
      • data_processing.py
      • environment.py
      • reporting.py
    • verify_dataset_weights.py
  • tests
    • conftest.py
    • run_all_tests.py
    • test_adaptive_logic.py
    • test_checkpointing.py
    • test_chinchilla.py
    • test_dataset.py
    • test_diagnose.py
    • test_e2e.py
    • test_initialization.py
    • test_integration.py
    • test_model.py
    • test_orchestrator.py
    • test_performance.py
    • test_tokenizer.py
    • test_trainer.py
    • test_unified_ops.py