Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .main.commit
Original file line number Diff line number Diff line change
@@ -1 +1 @@
8949def81eb744b9219e2d794c80265928c40a5b
fde43b6d617752433310b4668c6b5c55e503521a
2 changes: 1 addition & 1 deletion 3rdparty/Megatron-LM
Submodule Megatron-LM updated 65 files
+11 −0 .github/workflows/cicd-main.yml
+4 −0 .github/workflows/process-group-usage-check.yml
+3 −0 docker/Dockerfile.ci.dev
+37 −27 examples/rl/benchmark_refit.py
+1 −1 examples/rl/model_configs/common.sh
+143 −0 megatron/core/context_parallel/sequence_roll.py
+19 −3 megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/layout.py
+10 −10 megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/owner_planning.py
+8 −14 megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/quantized_dbuffer.py
+94 −0 megatron/core/fusions/fused_row_copy.py
+2 −1 megatron/core/inference/contexts/dynamic_context.py
+1 −2 megatron/core/inference/engines/dynamic_engine.py
+2 −2 megatron/core/inference/text_generation_controllers/mtp_controller_mixin.py
+6 −12 megatron/core/inference/text_generation_controllers/text_generation_controller.py
+11 −0 megatron/core/inference/utils.py
+1 −1 megatron/core/models/common/utils.py
+1 −1 megatron/core/models/gpt/fine_grained_callables.py
+95 −0 megatron/core/tensor_parallel/mappings.py
+509 −2 megatron/core/transformer/experimental_attention_variant/csa.py
+372 −0 megatron/core/transformer/experimental_attention_variant/csa_utils/cp_utils.py
+72 −0 megatron/core/transformer/experimental_attention_variant/csa_utils/csa_indexer_loss_kernels.py
+199 −35 megatron/core/transformer/experimental_attention_variant/csa_utils/csa_teacher_lse.py
+252 −0 megatron/core/transformer/experimental_attention_variant/csa_utils/packed_layout.py
+727 −0 megatron/core/transformer/experimental_attention_variant/csa_utils/packed_sparse_attention.py
+346 −115 megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py
+1 −0 megatron/core/transformer/moe/README.md
+232 −1 megatron/core/transformer/moe/fused_a2a.py
+2 −2 megatron/core/transformer/moe/moe_utils.py
+124 −1 megatron/core/transformer/moe/token_dispatcher.py
+17 −2 megatron/core/transformer/multi_latent_attention.py
+125 −70 megatron/core/transformer/multi_token_prediction.py
+20 −12 megatron/core/transformer/transformer_config.py
+1 −8 megatron/inference/utils.py
+58 −40 megatron/rl/inference/megatron.py
+53 −46 megatron/rl/rl_utils.py
+26 −85 megatron/training/arguments.py
+5 −2 megatron/training/checkpointing.py
+15 −0 megatron/training/config/container.py
+51 −18 megatron/training/config/inference_config.py
+55 −15 megatron/training/config/rl_config.py
+67 −59 megatron/training/training.py
+2 −2 tests/functional_tests/dynamic_inference_functional_tests.md
+2 −2 .../test_cases/moe/gpt_dynamic_inference_tp4_etp1_pp1_ep8_16B_logitsmatch_zmq_suspend_resume/model_config.yaml
+29 −1 tests/unit_tests/determinism/kernels/manifest.py
+100 −0 tests/unit_tests/determinism/kernels/test_csa_packed_kernels.py
+95 −1 tests/unit_tests/determinism/kernels/test_fused_triton_kernels.py
+22 −4 tests/unit_tests/determinism/kernels/test_moe_kernels.py
+36 −32 tests/unit_tests/distributed/mfsdp_v2/test_owner_planning.py
+42 −0 tests/unit_tests/fusions/test_fused_row_copy.py
+102 −21 tests/unit_tests/inference/test_inference_config.py
+24 −1 tests/unit_tests/inference/test_inference_utils.py
+1 −0 tests/unit_tests/inference/text_generation_controllers/test_mtp_kv_cache_controller.py
+24 −11 tests/unit_tests/rl/test_rl_utils.py
+1 −1 tests/unit_tests/rl/test_rollout_generation.py
+0 −12 tests/unit_tests/test_check_process_group_usage.py
+0 −1 tests/unit_tests/test_eval_batch_size.py
+2 −3 tests/unit_tests/test_training.py
+104 −8 tests/unit_tests/training/config/test_rl_config.py
+5 −3 tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention.py
+346 −0 tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_packed_cp.py
+13 −8 tests/unit_tests/transformer/moe/test_grouped_tensor_dispatcher_numerics.py
+13 −3 tests/unit_tests/transformer/moe/test_token_dispatcher.py
+67 −0 tests/unit_tests/transformer/test_contiguous_sequence_roll.py
+64 −0 tests/unit_tests/transformer/test_multi_latent_attention.py
+12 −12 train_rl.py
Original file line number Diff line number Diff line change
Expand Up @@ -254,6 +254,8 @@ def fake_output_layer(hidden, weight=None, runtime_gather_output=None): # noqa:
mtp_loss_scaling_factor=1.0,
calculate_per_token_loss=True,
use_mup=False,
# process_mtp_loss reads the CP layout; mirror TransformerConfig's default.
attention_cp_layout="zigzag",
)
output = process_mtp_loss(
hidden_states=hidden_states,
Expand Down
Loading
Loading