From 43eacb6181d49ea129708467398357e4ba01494b Mon Sep 17 00:00:00 2001 From: vidushi8 Date: Thu, 12 Feb 2026 19:13:55 -0600 Subject: [PATCH 01/11] Update torchtitan batxh size and enable CE fusion --- .../MI300X/deepseek_v2-BF16-pretrain.yaml | 4 ++-- .../MI300X/deepseek_v2-FP8-pretrain.yaml | 4 ++-- .../MI300X/deepseek_v2_lite-BF16-pretrain.yaml | 4 ++-- .../MI300X/deepseek_v2_lite-FP8-pretrain.yaml | 4 ++-- .../MI300X/deepseek_v3-BF16-pretrain.yaml | 4 ++-- .../MI300X/deepseek_v3-FP8-pretrain.yaml | 4 ++-- .../MI300X/llama2_70B-BF16-pretrain.yaml | 4 ++-- .../configs/MI300X/llama2_70B-FP8-pretrain.yaml | 4 ++-- .../configs/MI300X/llama2_7B-BF16-pretrain.yaml | 4 ++-- .../configs/MI300X/llama2_7B-FP8-pretrain.yaml | 4 ++-- .../MI300X/llama3.1_70B-BF16-pretrain.yaml | 4 ++-- .../MI300X/llama3.1_70B-FP8-pretrain.yaml | 4 ++-- .../MI300X/llama3.1_8B-BF16-pretrain.yaml | 4 ++-- .../MI300X/llama3.1_8B-FP8-pretrain.yaml | 4 ++-- .../MI300X/llama3_70B-BF16-pretrain.yaml | 4 ++-- .../configs/MI300X/llama3_70B-FP8-pretrain.yaml | 4 ++-- .../configs/MI300X/llama3_8B-BF16-pretrain.yaml | 4 ++-- .../configs/MI300X/llama3_8B-FP8-pretrain.yaml | 4 ++-- .../mixtral_8x22B_v0.1-BF16-pretrain.yaml | 4 ++-- .../MI300X/mixtral_8x22B_v0.1-FP8-pretrain.yaml | 4 ++-- .../MI300X/mixtral_8x7B_v0.1-BF16-pretrain.yaml | 4 ++-- .../MI300X/mixtral_8x7B_v0.1-FP8-pretrain.yaml | 4 ++-- .../MI300X/qwen2.5_72B-BF16-pretrain.yaml | 4 ++-- .../MI300X/qwen2.5_72B-FP8-pretrain.yaml | 4 ++-- .../MI300X/qwen2.5_7B-BF16-pretrain.yaml | 4 ++-- .../configs/MI300X/qwen2.5_7B-FP8-pretrain.yaml | 4 ++-- .../MI355X/deepseek_v2-BF16-pretrain.yaml | 4 ++-- .../MI355X/deepseek_v2-FP8-pretrain.yaml | 4 ++-- .../MI355X/deepseek_v2_lite-BF16-pretrain.yaml | 4 ++-- .../MI355X/deepseek_v2_lite-FP8-pretrain.yaml | 4 ++-- .../MI355X/deepseek_v3-BF16-pretrain.yaml | 17 +++++++++++++++-- .../MI355X/deepseek_v3-FP8-pretrain.yaml | 17 +++++++++++++++-- .../MI355X/llama2_70B-BF16-pretrain.yaml | 4 ++-- .../configs/MI355X/llama2_70B-FP8-pretrain.yaml | 4 ++-- .../configs/MI355X/llama2_7B-BF16-pretrain.yaml | 4 ++-- .../configs/MI355X/llama2_7B-FP8-pretrain.yaml | 4 ++-- .../MI355X/llama3.1_70B-BF16-pretrain.yaml | 4 ++-- .../MI355X/llama3.1_70B-FP8-pretrain.yaml | 4 ++-- .../MI355X/llama3.1_8B-BF16-pretrain.yaml | 4 ++-- .../MI355X/llama3.1_8B-FP8-pretrain.yaml | 4 ++-- .../MI355X/llama3.3_70B-BF16-pretrain.yaml | 4 ++-- .../MI355X/llama3.3_70B-FP8-pretrain.yaml | 4 ++-- .../MI355X/llama3_70B-BF16-pretrain.yaml | 4 ++-- .../configs/MI355X/llama3_70B-FP8-pretrain.yaml | 4 ++-- .../configs/MI355X/llama3_8B-BF16-pretrain.yaml | 4 ++-- .../configs/MI355X/llama3_8B-FP8-pretrain.yaml | 4 ++-- .../mixtral_8x22B_v0.1-BF16-pretrain.yaml | 4 ++-- .../MI355X/mixtral_8x22B_v0.1-FP8-pretrain.yaml | 4 ++-- .../MI355X/mixtral_8x7B_v0.1-BF16-pretrain.yaml | 4 ++-- .../MI355X/mixtral_8x7B_v0.1-FP8-pretrain.yaml | 4 ++-- .../MI355X/qwen2.5_72B-BF16-pretrain.yaml | 4 ++-- .../MI355X/qwen2.5_72B-FP8-pretrain.yaml | 4 ++-- .../MI355X/qwen2.5_7B-BF16-pretrain.yaml | 4 ++-- .../configs/MI355X/qwen2.5_7B-FP8-pretrain.yaml | 4 ++-- .../MI300X/llama3.1_70B-BF16-pretrain.yaml | 2 +- .../MI300X/llama3.1_70B-FP8-pretrain.yaml | 2 +- .../MI355X/llama3.1_70B-BF16-pretrain.yaml | 2 +- .../MI355X/llama3.1_70B-FP8-pretrain.yaml | 2 +- 58 files changed, 138 insertions(+), 112 deletions(-) diff --git a/examples/megatron/configs/MI300X/deepseek_v2-BF16-pretrain.yaml b/examples/megatron/configs/MI300X/deepseek_v2-BF16-pretrain.yaml index 24f042f4c..234ad2953 100644 --- a/examples/megatron/configs/MI300X/deepseek_v2-BF16-pretrain.yaml +++ b/examples/megatron/configs/MI300X/deepseek_v2-BF16-pretrain.yaml @@ -95,5 +95,5 @@ modules: turbo_sync_free_moe_stage: 2 # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true diff --git a/examples/megatron/configs/MI300X/deepseek_v2-FP8-pretrain.yaml b/examples/megatron/configs/MI300X/deepseek_v2-FP8-pretrain.yaml index 8f4d0570b..42dcc754f 100644 --- a/examples/megatron/configs/MI300X/deepseek_v2-FP8-pretrain.yaml +++ b/examples/megatron/configs/MI300X/deepseek_v2-FP8-pretrain.yaml @@ -94,8 +94,8 @@ modules: turbo_sync_free_moe_stage: 2 # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true # enable fp8 training fp8: hybrid diff --git a/examples/megatron/configs/MI300X/deepseek_v2_lite-BF16-pretrain.yaml b/examples/megatron/configs/MI300X/deepseek_v2_lite-BF16-pretrain.yaml index a43928a3b..d97c22d1e 100644 --- a/examples/megatron/configs/MI300X/deepseek_v2_lite-BF16-pretrain.yaml +++ b/examples/megatron/configs/MI300X/deepseek_v2_lite-BF16-pretrain.yaml @@ -90,5 +90,5 @@ modules: turbo_sync_free_moe_stage: 2 # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true diff --git a/examples/megatron/configs/MI300X/deepseek_v2_lite-FP8-pretrain.yaml b/examples/megatron/configs/MI300X/deepseek_v2_lite-FP8-pretrain.yaml index 794a305e2..950fe9154 100644 --- a/examples/megatron/configs/MI300X/deepseek_v2_lite-FP8-pretrain.yaml +++ b/examples/megatron/configs/MI300X/deepseek_v2_lite-FP8-pretrain.yaml @@ -89,8 +89,8 @@ modules: turbo_sync_free_moe_stage: 2 # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true # enable fp8 training fp8: hybrid diff --git a/examples/megatron/configs/MI300X/deepseek_v3-BF16-pretrain.yaml b/examples/megatron/configs/MI300X/deepseek_v3-BF16-pretrain.yaml index 825f131e0..df75699f3 100644 --- a/examples/megatron/configs/MI300X/deepseek_v3-BF16-pretrain.yaml +++ b/examples/megatron/configs/MI300X/deepseek_v3-BF16-pretrain.yaml @@ -80,5 +80,5 @@ modules: use_turbo_grouped_mlp: true # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true diff --git a/examples/megatron/configs/MI300X/deepseek_v3-FP8-pretrain.yaml b/examples/megatron/configs/MI300X/deepseek_v3-FP8-pretrain.yaml index b7cfe474f..68ca7064a 100644 --- a/examples/megatron/configs/MI300X/deepseek_v3-FP8-pretrain.yaml +++ b/examples/megatron/configs/MI300X/deepseek_v3-FP8-pretrain.yaml @@ -79,8 +79,8 @@ modules: use_turbo_grouped_mlp: true # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true # enable fp8 training fp8: hybrid diff --git a/examples/megatron/configs/MI300X/llama2_70B-BF16-pretrain.yaml b/examples/megatron/configs/MI300X/llama2_70B-BF16-pretrain.yaml index 1493fd25f..e0f48d9a4 100755 --- a/examples/megatron/configs/MI300X/llama2_70B-BF16-pretrain.yaml +++ b/examples/megatron/configs/MI300X/llama2_70B-BF16-pretrain.yaml @@ -77,5 +77,5 @@ modules: use_turbo_grouped_mlp: true # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true diff --git a/examples/megatron/configs/MI300X/llama2_70B-FP8-pretrain.yaml b/examples/megatron/configs/MI300X/llama2_70B-FP8-pretrain.yaml index 8f83cfd8d..77d4b52bc 100755 --- a/examples/megatron/configs/MI300X/llama2_70B-FP8-pretrain.yaml +++ b/examples/megatron/configs/MI300X/llama2_70B-FP8-pretrain.yaml @@ -77,8 +77,8 @@ modules: use_turbo_grouped_mlp: true # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true # enable fp8 training fp8: hybrid diff --git a/examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml b/examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml index 19cb8465f..0ceeeae27 100755 --- a/examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml +++ b/examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml @@ -80,5 +80,5 @@ modules: # sequence_parallel: 1 # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true diff --git a/examples/megatron/configs/MI300X/llama2_7B-FP8-pretrain.yaml b/examples/megatron/configs/MI300X/llama2_7B-FP8-pretrain.yaml index 73a79b00b..347a3430a 100755 --- a/examples/megatron/configs/MI300X/llama2_7B-FP8-pretrain.yaml +++ b/examples/megatron/configs/MI300X/llama2_7B-FP8-pretrain.yaml @@ -80,8 +80,8 @@ modules: # sequence_parallel: 1 # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true # enable fp8 training fp8: hybrid diff --git a/examples/megatron/configs/MI300X/llama3.1_70B-BF16-pretrain.yaml b/examples/megatron/configs/MI300X/llama3.1_70B-BF16-pretrain.yaml index bb2454f1e..5b67458b8 100644 --- a/examples/megatron/configs/MI300X/llama3.1_70B-BF16-pretrain.yaml +++ b/examples/megatron/configs/MI300X/llama3.1_70B-BF16-pretrain.yaml @@ -76,5 +76,5 @@ modules: use_turbo_attention: true # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true diff --git a/examples/megatron/configs/MI300X/llama3.1_70B-FP8-pretrain.yaml b/examples/megatron/configs/MI300X/llama3.1_70B-FP8-pretrain.yaml index 9013b0c96..a94a66c2e 100644 --- a/examples/megatron/configs/MI300X/llama3.1_70B-FP8-pretrain.yaml +++ b/examples/megatron/configs/MI300X/llama3.1_70B-FP8-pretrain.yaml @@ -76,8 +76,8 @@ modules: use_turbo_attention: true # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true # enable fp8 training fp8: hybrid diff --git a/examples/megatron/configs/MI300X/llama3.1_8B-BF16-pretrain.yaml b/examples/megatron/configs/MI300X/llama3.1_8B-BF16-pretrain.yaml index 4d49c5c09..6a4976ada 100644 --- a/examples/megatron/configs/MI300X/llama3.1_8B-BF16-pretrain.yaml +++ b/examples/megatron/configs/MI300X/llama3.1_8B-BF16-pretrain.yaml @@ -72,5 +72,5 @@ modules: use_turbo_grouped_mlp: false # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true diff --git a/examples/megatron/configs/MI300X/llama3.1_8B-FP8-pretrain.yaml b/examples/megatron/configs/MI300X/llama3.1_8B-FP8-pretrain.yaml index 9e4b65560..148be586a 100644 --- a/examples/megatron/configs/MI300X/llama3.1_8B-FP8-pretrain.yaml +++ b/examples/megatron/configs/MI300X/llama3.1_8B-FP8-pretrain.yaml @@ -72,8 +72,8 @@ modules: use_turbo_grouped_mlp: false # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true # enable fp8 training fp8: hybrid diff --git a/examples/megatron/configs/MI300X/llama3_70B-BF16-pretrain.yaml b/examples/megatron/configs/MI300X/llama3_70B-BF16-pretrain.yaml index 468c56827..68d035a2b 100755 --- a/examples/megatron/configs/MI300X/llama3_70B-BF16-pretrain.yaml +++ b/examples/megatron/configs/MI300X/llama3_70B-BF16-pretrain.yaml @@ -77,5 +77,5 @@ modules: use_turbo_grouped_mlp: true # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true diff --git a/examples/megatron/configs/MI300X/llama3_70B-FP8-pretrain.yaml b/examples/megatron/configs/MI300X/llama3_70B-FP8-pretrain.yaml index 8e4ee47a2..23c157f62 100755 --- a/examples/megatron/configs/MI300X/llama3_70B-FP8-pretrain.yaml +++ b/examples/megatron/configs/MI300X/llama3_70B-FP8-pretrain.yaml @@ -77,8 +77,8 @@ modules: use_turbo_grouped_mlp: true # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true # enable fp8 training fp8: hybrid diff --git a/examples/megatron/configs/MI300X/llama3_8B-BF16-pretrain.yaml b/examples/megatron/configs/MI300X/llama3_8B-BF16-pretrain.yaml index 7d2588f95..6af89e06b 100644 --- a/examples/megatron/configs/MI300X/llama3_8B-BF16-pretrain.yaml +++ b/examples/megatron/configs/MI300X/llama3_8B-BF16-pretrain.yaml @@ -73,5 +73,5 @@ modules: use_turbo_grouped_mlp: true # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true diff --git a/examples/megatron/configs/MI300X/llama3_8B-FP8-pretrain.yaml b/examples/megatron/configs/MI300X/llama3_8B-FP8-pretrain.yaml index ef9149801..0ad398267 100644 --- a/examples/megatron/configs/MI300X/llama3_8B-FP8-pretrain.yaml +++ b/examples/megatron/configs/MI300X/llama3_8B-FP8-pretrain.yaml @@ -73,8 +73,8 @@ modules: use_turbo_grouped_mlp: true # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true # enable fp8 training fp8: hybrid diff --git a/examples/megatron/configs/MI300X/mixtral_8x22B_v0.1-BF16-pretrain.yaml b/examples/megatron/configs/MI300X/mixtral_8x22B_v0.1-BF16-pretrain.yaml index 0efc4ee0b..73e949215 100644 --- a/examples/megatron/configs/MI300X/mixtral_8x22B_v0.1-BF16-pretrain.yaml +++ b/examples/megatron/configs/MI300X/mixtral_8x22B_v0.1-BF16-pretrain.yaml @@ -73,5 +73,5 @@ modules: ckpt_format: torch # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true diff --git a/examples/megatron/configs/MI300X/mixtral_8x22B_v0.1-FP8-pretrain.yaml b/examples/megatron/configs/MI300X/mixtral_8x22B_v0.1-FP8-pretrain.yaml index 520177a76..14a554c5e 100644 --- a/examples/megatron/configs/MI300X/mixtral_8x22B_v0.1-FP8-pretrain.yaml +++ b/examples/megatron/configs/MI300X/mixtral_8x22B_v0.1-FP8-pretrain.yaml @@ -72,8 +72,8 @@ modules: ckpt_format: torch # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true # enable fp8 training fp8: hybrid diff --git a/examples/megatron/configs/MI300X/mixtral_8x7B_v0.1-BF16-pretrain.yaml b/examples/megatron/configs/MI300X/mixtral_8x7B_v0.1-BF16-pretrain.yaml index eb5ce376f..28a41749e 100644 --- a/examples/megatron/configs/MI300X/mixtral_8x7B_v0.1-BF16-pretrain.yaml +++ b/examples/megatron/configs/MI300X/mixtral_8x7B_v0.1-BF16-pretrain.yaml @@ -68,5 +68,5 @@ modules: ckpt_format: torch # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true diff --git a/examples/megatron/configs/MI300X/mixtral_8x7B_v0.1-FP8-pretrain.yaml b/examples/megatron/configs/MI300X/mixtral_8x7B_v0.1-FP8-pretrain.yaml index 35fa84415..c92eea95f 100644 --- a/examples/megatron/configs/MI300X/mixtral_8x7B_v0.1-FP8-pretrain.yaml +++ b/examples/megatron/configs/MI300X/mixtral_8x7B_v0.1-FP8-pretrain.yaml @@ -67,8 +67,8 @@ modules: ckpt_format: torch # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true # enable fp8 training fp8: hybrid diff --git a/examples/megatron/configs/MI300X/qwen2.5_72B-BF16-pretrain.yaml b/examples/megatron/configs/MI300X/qwen2.5_72B-BF16-pretrain.yaml index e02c3a233..ab300f580 100644 --- a/examples/megatron/configs/MI300X/qwen2.5_72B-BF16-pretrain.yaml +++ b/examples/megatron/configs/MI300X/qwen2.5_72B-BF16-pretrain.yaml @@ -81,5 +81,5 @@ modules: use_turbo_grouped_mlp: true # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true diff --git a/examples/megatron/configs/MI300X/qwen2.5_72B-FP8-pretrain.yaml b/examples/megatron/configs/MI300X/qwen2.5_72B-FP8-pretrain.yaml index 43f19a513..84317c0d5 100644 --- a/examples/megatron/configs/MI300X/qwen2.5_72B-FP8-pretrain.yaml +++ b/examples/megatron/configs/MI300X/qwen2.5_72B-FP8-pretrain.yaml @@ -81,8 +81,8 @@ modules: use_turbo_grouped_mlp: true # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true # enable fp8 training fp8: hybrid diff --git a/examples/megatron/configs/MI300X/qwen2.5_7B-BF16-pretrain.yaml b/examples/megatron/configs/MI300X/qwen2.5_7B-BF16-pretrain.yaml index 8ab17cbdd..b9d9fbe63 100644 --- a/examples/megatron/configs/MI300X/qwen2.5_7B-BF16-pretrain.yaml +++ b/examples/megatron/configs/MI300X/qwen2.5_7B-BF16-pretrain.yaml @@ -74,5 +74,5 @@ modules: use_turbo_grouped_mlp: true # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true diff --git a/examples/megatron/configs/MI300X/qwen2.5_7B-FP8-pretrain.yaml b/examples/megatron/configs/MI300X/qwen2.5_7B-FP8-pretrain.yaml index 4f331233e..93c2fcf68 100644 --- a/examples/megatron/configs/MI300X/qwen2.5_7B-FP8-pretrain.yaml +++ b/examples/megatron/configs/MI300X/qwen2.5_7B-FP8-pretrain.yaml @@ -74,8 +74,8 @@ modules: use_turbo_grouped_mlp: true # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true # enable fp8 training fp8: hybrid diff --git a/examples/megatron/configs/MI355X/deepseek_v2-BF16-pretrain.yaml b/examples/megatron/configs/MI355X/deepseek_v2-BF16-pretrain.yaml index 987b0d019..5f7d1299c 100644 --- a/examples/megatron/configs/MI355X/deepseek_v2-BF16-pretrain.yaml +++ b/examples/megatron/configs/MI355X/deepseek_v2-BF16-pretrain.yaml @@ -94,5 +94,5 @@ modules: turbo_sync_free_moe_stage: 1 # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true diff --git a/examples/megatron/configs/MI355X/deepseek_v2-FP8-pretrain.yaml b/examples/megatron/configs/MI355X/deepseek_v2-FP8-pretrain.yaml index 328eb3216..8c91eabb0 100644 --- a/examples/megatron/configs/MI355X/deepseek_v2-FP8-pretrain.yaml +++ b/examples/megatron/configs/MI355X/deepseek_v2-FP8-pretrain.yaml @@ -93,8 +93,8 @@ modules: turbo_sync_free_moe_stage: 1 # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true # enable fp8 training fp8: hybrid diff --git a/examples/megatron/configs/MI355X/deepseek_v2_lite-BF16-pretrain.yaml b/examples/megatron/configs/MI355X/deepseek_v2_lite-BF16-pretrain.yaml index b706c159f..cf8064671 100644 --- a/examples/megatron/configs/MI355X/deepseek_v2_lite-BF16-pretrain.yaml +++ b/examples/megatron/configs/MI355X/deepseek_v2_lite-BF16-pretrain.yaml @@ -89,5 +89,5 @@ modules: turbo_sync_free_moe_stage: 1 # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true diff --git a/examples/megatron/configs/MI355X/deepseek_v2_lite-FP8-pretrain.yaml b/examples/megatron/configs/MI355X/deepseek_v2_lite-FP8-pretrain.yaml index 4f3a9f679..715d05551 100644 --- a/examples/megatron/configs/MI355X/deepseek_v2_lite-FP8-pretrain.yaml +++ b/examples/megatron/configs/MI355X/deepseek_v2_lite-FP8-pretrain.yaml @@ -88,8 +88,8 @@ modules: turbo_sync_free_moe_stage: 1 # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true # enable fp8 training fp8: hybrid diff --git a/examples/megatron/configs/MI355X/deepseek_v3-BF16-pretrain.yaml b/examples/megatron/configs/MI355X/deepseek_v3-BF16-pretrain.yaml index 51dd8affa..61bd8aa5d 100644 --- a/examples/megatron/configs/MI355X/deepseek_v3-BF16-pretrain.yaml +++ b/examples/megatron/configs/MI355X/deepseek_v3-BF16-pretrain.yaml @@ -71,6 +71,19 @@ modules: ckpt_format: torch eval_iters: 0 + # deepep + use_turbo_deepep: true + moe_shared_expert_overlap: false + moe_router_dtype: fp32 + + # 64 or 80 for ep8, 32 for ep16-64 is best practice + turbo_deepep_num_cu: 64 + turbo_deepep_use_comm_stream: false + + # sync-free moe support stage 1-2, 0 means not use sync-free moe + # stage 2 is recommended for better performance + turbo_sync_free_moe_stage: 1 + # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true diff --git a/examples/megatron/configs/MI355X/deepseek_v3-FP8-pretrain.yaml b/examples/megatron/configs/MI355X/deepseek_v3-FP8-pretrain.yaml index a40c24449..08e109c7b 100644 --- a/examples/megatron/configs/MI355X/deepseek_v3-FP8-pretrain.yaml +++ b/examples/megatron/configs/MI355X/deepseek_v3-FP8-pretrain.yaml @@ -70,9 +70,22 @@ modules: ckpt_format: torch eval_iters: 0 + # deepep + use_turbo_deepep: true + moe_shared_expert_overlap: false + moe_router_dtype: fp32 + + # 64 or 80 for ep8, 32 for ep16-64 is best practice + turbo_deepep_num_cu: 64 + turbo_deepep_use_comm_stream: false + + # sync-free moe support stage 1-2, 0 means not use sync-free moe + # stage 2 is recommended for better performance + turbo_sync_free_moe_stage: 1 + # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true # enable fp8 training fp8: hybrid diff --git a/examples/megatron/configs/MI355X/llama2_70B-BF16-pretrain.yaml b/examples/megatron/configs/MI355X/llama2_70B-BF16-pretrain.yaml index 87b933fdd..2d07839f9 100755 --- a/examples/megatron/configs/MI355X/llama2_70B-BF16-pretrain.yaml +++ b/examples/megatron/configs/MI355X/llama2_70B-BF16-pretrain.yaml @@ -77,5 +77,5 @@ modules: use_turbo_grouped_mlp: false # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true diff --git a/examples/megatron/configs/MI355X/llama2_70B-FP8-pretrain.yaml b/examples/megatron/configs/MI355X/llama2_70B-FP8-pretrain.yaml index 903bd2ad6..923c872df 100755 --- a/examples/megatron/configs/MI355X/llama2_70B-FP8-pretrain.yaml +++ b/examples/megatron/configs/MI355X/llama2_70B-FP8-pretrain.yaml @@ -77,8 +77,8 @@ modules: use_turbo_grouped_mlp: false # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true # enable fp8 training fp8: hybrid diff --git a/examples/megatron/configs/MI355X/llama2_7B-BF16-pretrain.yaml b/examples/megatron/configs/MI355X/llama2_7B-BF16-pretrain.yaml index 641dbbaf1..607fe5477 100755 --- a/examples/megatron/configs/MI355X/llama2_7B-BF16-pretrain.yaml +++ b/examples/megatron/configs/MI355X/llama2_7B-BF16-pretrain.yaml @@ -80,5 +80,5 @@ modules: # sequence_parallel: 1 # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true diff --git a/examples/megatron/configs/MI355X/llama2_7B-FP8-pretrain.yaml b/examples/megatron/configs/MI355X/llama2_7B-FP8-pretrain.yaml index 43548e58e..b77adb414 100755 --- a/examples/megatron/configs/MI355X/llama2_7B-FP8-pretrain.yaml +++ b/examples/megatron/configs/MI355X/llama2_7B-FP8-pretrain.yaml @@ -80,8 +80,8 @@ modules: # sequence_parallel: 1 # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true # enable fp8 training fp8: hybrid diff --git a/examples/megatron/configs/MI355X/llama3.1_70B-BF16-pretrain.yaml b/examples/megatron/configs/MI355X/llama3.1_70B-BF16-pretrain.yaml index f03d9c4c1..6ee1f25b3 100644 --- a/examples/megatron/configs/MI355X/llama3.1_70B-BF16-pretrain.yaml +++ b/examples/megatron/configs/MI355X/llama3.1_70B-BF16-pretrain.yaml @@ -72,5 +72,5 @@ modules: recompute_num_layers: 80 # int # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true diff --git a/examples/megatron/configs/MI355X/llama3.1_70B-FP8-pretrain.yaml b/examples/megatron/configs/MI355X/llama3.1_70B-FP8-pretrain.yaml index ae82330d5..2c2fa0b68 100644 --- a/examples/megatron/configs/MI355X/llama3.1_70B-FP8-pretrain.yaml +++ b/examples/megatron/configs/MI355X/llama3.1_70B-FP8-pretrain.yaml @@ -72,8 +72,8 @@ modules: recompute_num_layers: 80 # int # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true # enable fp8 training fp8: hybrid diff --git a/examples/megatron/configs/MI355X/llama3.1_8B-BF16-pretrain.yaml b/examples/megatron/configs/MI355X/llama3.1_8B-BF16-pretrain.yaml index 59a7f5b8d..83722456d 100644 --- a/examples/megatron/configs/MI355X/llama3.1_8B-BF16-pretrain.yaml +++ b/examples/megatron/configs/MI355X/llama3.1_8B-BF16-pretrain.yaml @@ -67,5 +67,5 @@ modules: ckpt_format: torch # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true diff --git a/examples/megatron/configs/MI355X/llama3.1_8B-FP8-pretrain.yaml b/examples/megatron/configs/MI355X/llama3.1_8B-FP8-pretrain.yaml index 5cba8ee79..cff682659 100644 --- a/examples/megatron/configs/MI355X/llama3.1_8B-FP8-pretrain.yaml +++ b/examples/megatron/configs/MI355X/llama3.1_8B-FP8-pretrain.yaml @@ -67,8 +67,8 @@ modules: ckpt_format: torch # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true # enable fp8 training fp8: hybrid diff --git a/examples/megatron/configs/MI355X/llama3.3_70B-BF16-pretrain.yaml b/examples/megatron/configs/MI355X/llama3.3_70B-BF16-pretrain.yaml index 3316fbd4d..fa860a013 100644 --- a/examples/megatron/configs/MI355X/llama3.3_70B-BF16-pretrain.yaml +++ b/examples/megatron/configs/MI355X/llama3.3_70B-BF16-pretrain.yaml @@ -77,5 +77,5 @@ modules: use_turbo_grouped_mlp: false # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true diff --git a/examples/megatron/configs/MI355X/llama3.3_70B-FP8-pretrain.yaml b/examples/megatron/configs/MI355X/llama3.3_70B-FP8-pretrain.yaml index f79ef9bbd..a7727049e 100644 --- a/examples/megatron/configs/MI355X/llama3.3_70B-FP8-pretrain.yaml +++ b/examples/megatron/configs/MI355X/llama3.3_70B-FP8-pretrain.yaml @@ -77,8 +77,8 @@ modules: use_turbo_grouped_mlp: false # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true # enable fp8 training fp8: hybrid diff --git a/examples/megatron/configs/MI355X/llama3_70B-BF16-pretrain.yaml b/examples/megatron/configs/MI355X/llama3_70B-BF16-pretrain.yaml index ca25b05d2..00554eee9 100755 --- a/examples/megatron/configs/MI355X/llama3_70B-BF16-pretrain.yaml +++ b/examples/megatron/configs/MI355X/llama3_70B-BF16-pretrain.yaml @@ -77,5 +77,5 @@ modules: use_turbo_grouped_mlp: false # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true diff --git a/examples/megatron/configs/MI355X/llama3_70B-FP8-pretrain.yaml b/examples/megatron/configs/MI355X/llama3_70B-FP8-pretrain.yaml index 22bb90699..3de4437eb 100755 --- a/examples/megatron/configs/MI355X/llama3_70B-FP8-pretrain.yaml +++ b/examples/megatron/configs/MI355X/llama3_70B-FP8-pretrain.yaml @@ -77,8 +77,8 @@ modules: use_turbo_grouped_mlp: false # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true # enable fp8 training fp8: hybrid diff --git a/examples/megatron/configs/MI355X/llama3_8B-BF16-pretrain.yaml b/examples/megatron/configs/MI355X/llama3_8B-BF16-pretrain.yaml index 2093e168e..7bfc7b1bb 100644 --- a/examples/megatron/configs/MI355X/llama3_8B-BF16-pretrain.yaml +++ b/examples/megatron/configs/MI355X/llama3_8B-BF16-pretrain.yaml @@ -73,5 +73,5 @@ modules: use_turbo_grouped_mlp: false # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true diff --git a/examples/megatron/configs/MI355X/llama3_8B-FP8-pretrain.yaml b/examples/megatron/configs/MI355X/llama3_8B-FP8-pretrain.yaml index 1ec470ad5..096345c92 100644 --- a/examples/megatron/configs/MI355X/llama3_8B-FP8-pretrain.yaml +++ b/examples/megatron/configs/MI355X/llama3_8B-FP8-pretrain.yaml @@ -73,8 +73,8 @@ modules: use_turbo_grouped_mlp: false # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true # enable fp8 training fp8: hybrid diff --git a/examples/megatron/configs/MI355X/mixtral_8x22B_v0.1-BF16-pretrain.yaml b/examples/megatron/configs/MI355X/mixtral_8x22B_v0.1-BF16-pretrain.yaml index 188465761..0b30cadf5 100644 --- a/examples/megatron/configs/MI355X/mixtral_8x22B_v0.1-BF16-pretrain.yaml +++ b/examples/megatron/configs/MI355X/mixtral_8x22B_v0.1-BF16-pretrain.yaml @@ -73,8 +73,8 @@ modules: ckpt_format: torch # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true # Turbo enable_primus_turbo: true diff --git a/examples/megatron/configs/MI355X/mixtral_8x22B_v0.1-FP8-pretrain.yaml b/examples/megatron/configs/MI355X/mixtral_8x22B_v0.1-FP8-pretrain.yaml index ea9beca43..712f4a022 100644 --- a/examples/megatron/configs/MI355X/mixtral_8x22B_v0.1-FP8-pretrain.yaml +++ b/examples/megatron/configs/MI355X/mixtral_8x22B_v0.1-FP8-pretrain.yaml @@ -72,8 +72,8 @@ modules: ckpt_format: torch # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true # Turbo enable_primus_turbo: true diff --git a/examples/megatron/configs/MI355X/mixtral_8x7B_v0.1-BF16-pretrain.yaml b/examples/megatron/configs/MI355X/mixtral_8x7B_v0.1-BF16-pretrain.yaml index 16d8ffbcb..709373f3d 100644 --- a/examples/megatron/configs/MI355X/mixtral_8x7B_v0.1-BF16-pretrain.yaml +++ b/examples/megatron/configs/MI355X/mixtral_8x7B_v0.1-BF16-pretrain.yaml @@ -68,8 +68,8 @@ modules: ckpt_format: torch # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true # Turbo enable_primus_turbo: true diff --git a/examples/megatron/configs/MI355X/mixtral_8x7B_v0.1-FP8-pretrain.yaml b/examples/megatron/configs/MI355X/mixtral_8x7B_v0.1-FP8-pretrain.yaml index 4951e3d66..ee2684c3c 100644 --- a/examples/megatron/configs/MI355X/mixtral_8x7B_v0.1-FP8-pretrain.yaml +++ b/examples/megatron/configs/MI355X/mixtral_8x7B_v0.1-FP8-pretrain.yaml @@ -67,8 +67,8 @@ modules: ckpt_format: torch # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true # Turbo enable_primus_turbo: true diff --git a/examples/megatron/configs/MI355X/qwen2.5_72B-BF16-pretrain.yaml b/examples/megatron/configs/MI355X/qwen2.5_72B-BF16-pretrain.yaml index aae072840..812617737 100644 --- a/examples/megatron/configs/MI355X/qwen2.5_72B-BF16-pretrain.yaml +++ b/examples/megatron/configs/MI355X/qwen2.5_72B-BF16-pretrain.yaml @@ -81,5 +81,5 @@ modules: use_turbo_grouped_mlp: false # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true diff --git a/examples/megatron/configs/MI355X/qwen2.5_72B-FP8-pretrain.yaml b/examples/megatron/configs/MI355X/qwen2.5_72B-FP8-pretrain.yaml index 8cf56be9e..747ff24cf 100644 --- a/examples/megatron/configs/MI355X/qwen2.5_72B-FP8-pretrain.yaml +++ b/examples/megatron/configs/MI355X/qwen2.5_72B-FP8-pretrain.yaml @@ -81,8 +81,8 @@ modules: use_turbo_grouped_mlp: false # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true # enable fp8 training fp8: hybrid diff --git a/examples/megatron/configs/MI355X/qwen2.5_7B-BF16-pretrain.yaml b/examples/megatron/configs/MI355X/qwen2.5_7B-BF16-pretrain.yaml index 0aa928e9b..f82a51992 100644 --- a/examples/megatron/configs/MI355X/qwen2.5_7B-BF16-pretrain.yaml +++ b/examples/megatron/configs/MI355X/qwen2.5_7B-BF16-pretrain.yaml @@ -74,5 +74,5 @@ modules: use_turbo_grouped_mlp: true # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true diff --git a/examples/megatron/configs/MI355X/qwen2.5_7B-FP8-pretrain.yaml b/examples/megatron/configs/MI355X/qwen2.5_7B-FP8-pretrain.yaml index 6e9f44515..872180b34 100644 --- a/examples/megatron/configs/MI355X/qwen2.5_7B-FP8-pretrain.yaml +++ b/examples/megatron/configs/MI355X/qwen2.5_7B-FP8-pretrain.yaml @@ -74,8 +74,8 @@ modules: use_turbo_grouped_mlp: true # Cross entropy flags - # cross_entropy_fusion_impl: "te" - # cross_entropy_loss_fusion: true + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true # enable fp8 training fp8: hybrid diff --git a/examples/torchtitan/configs/MI300X/llama3.1_70B-BF16-pretrain.yaml b/examples/torchtitan/configs/MI300X/llama3.1_70B-BF16-pretrain.yaml index da0831030..21970087a 100644 --- a/examples/torchtitan/configs/MI300X/llama3.1_70B-BF16-pretrain.yaml +++ b/examples/torchtitan/configs/MI300X/llama3.1_70B-BF16-pretrain.yaml @@ -23,7 +23,7 @@ modules: log_freq: 1 training: - local_batch_size: 4 + local_batch_size: 2 seq_len: 8192 mock_data: false steps: 50 diff --git a/examples/torchtitan/configs/MI300X/llama3.1_70B-FP8-pretrain.yaml b/examples/torchtitan/configs/MI300X/llama3.1_70B-FP8-pretrain.yaml index 93e0b458b..6e09f2dcf 100644 --- a/examples/torchtitan/configs/MI300X/llama3.1_70B-FP8-pretrain.yaml +++ b/examples/torchtitan/configs/MI300X/llama3.1_70B-FP8-pretrain.yaml @@ -23,7 +23,7 @@ modules: log_freq: 1 training: - local_batch_size: 3 + local_batch_size: 1 seq_len: 8192 mock_data: false steps: 50 diff --git a/examples/torchtitan/configs/MI355X/llama3.1_70B-BF16-pretrain.yaml b/examples/torchtitan/configs/MI355X/llama3.1_70B-BF16-pretrain.yaml index 23de79da8..1e4465d38 100644 --- a/examples/torchtitan/configs/MI355X/llama3.1_70B-BF16-pretrain.yaml +++ b/examples/torchtitan/configs/MI355X/llama3.1_70B-BF16-pretrain.yaml @@ -23,7 +23,7 @@ modules: log_freq: 1 training: - local_batch_size: 8 + local_batch_size: 3 seq_len: 8192 mock_data: false steps: 10 diff --git a/examples/torchtitan/configs/MI355X/llama3.1_70B-FP8-pretrain.yaml b/examples/torchtitan/configs/MI355X/llama3.1_70B-FP8-pretrain.yaml index 034d8523a..91b92ed32 100644 --- a/examples/torchtitan/configs/MI355X/llama3.1_70B-FP8-pretrain.yaml +++ b/examples/torchtitan/configs/MI355X/llama3.1_70B-FP8-pretrain.yaml @@ -28,7 +28,7 @@ modules: log_freq: 1 training: - local_batch_size: 6 + local_batch_size: 3 seq_len: 8192 mock_data: false steps: 10 From 33daa5ae394b4462d6236eaa2b833899443d508d Mon Sep 17 00:00:00 2001 From: Kailash Gogineni Date: Thu, 19 Feb 2026 07:21:32 -0800 Subject: [PATCH 02/11] [Docs] & [Feature]: Add Post-Training Documentation and Update Qwen3_32B Configs for MI300X & MI355X (#556) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit YF: Only SFT related config and Doc changes, bypassing unit CI tests ## Summary This PR introduces post-training documentation and updates Qwen3 32B model configuration files to support AMD MI300X and MI355X accelerators. --- ## Changes ### πŸ“˜ Documentation - **Added `posttraining.md`** - New comprehensive guide for post-training workflows - Covers setup instructions, configuration details, and usage examples - **Updated `docs/README.md`** - Added a new section referencing post-training documentation - Improved documentation organization and navigation --- ### βš™οΈ Configuration Updates - **Updated Qwen3_32B model YAML configs** - Added/modified configurations optimized for: - MI300X - MI355X - Adjusted parameters for compatibility and stable execution --- ## Validation - Verified updated configs load and execute successfully on MI300X and MI355X environments - Confirmed documentation links and structure render correctly --- ## Checklist - [x] Added `posttraining.md` - [x] Updated `docs/README.md` - [x] Modified Qwen3_32B YAML configs - [x] Verified changes locally --- docs/README.md | 1 + docs/posttraining.md | 452 ++++++++++++++++++ .../MI300X/qwen3_32b_lora_posttrain.yaml | 58 +++ .../MI300X/qwen3_32b_sft_posttrain.yaml | 58 +++ .../MI355X/qwen3_32b_lora_posttrain.yaml | 6 +- .../MI355X/qwen3_32b_sft_posttrain.yaml | 6 +- .../modules/megatron_bridge/sft_trainer.yaml | 2 +- 7 files changed, 576 insertions(+), 7 deletions(-) create mode 100644 docs/posttraining.md create mode 100644 examples/megatron_bridge/configs/MI300X/qwen3_32b_lora_posttrain.yaml create mode 100644 examples/megatron_bridge/configs/MI300X/qwen3_32b_sft_posttrain.yaml diff --git a/docs/README.md b/docs/README.md index 45e28f26e..f2d22ad42 100644 --- a/docs/README.md +++ b/docs/README.md @@ -24,6 +24,7 @@ Guides for common workflows and features: In-depth technical documentation: +- **[Post-Training Guide](./posttraining.md)** - Fine-tuning with SFT and LoRA using Primus CLI - **[Performance Projection](./projection.md)** - Project training performance to multi-node configurations - **[Preflight](./preflight.md)** - Cluster diagnostics (host/GPU/network info + perf tests) - **[Benchmark Suite](./benchmark.md)** - GEMM, RCCL, end-to-end benchmarks and profiling diff --git a/docs/posttraining.md b/docs/posttraining.md new file mode 100644 index 000000000..27eb5ca01 --- /dev/null +++ b/docs/posttraining.md @@ -0,0 +1,452 @@ +# πŸŽ“ Post-Training with Primus + +This guide demonstrates how to perform post-training (fine-tuning) using **Megatron Bridge** within the **Primus** framework. It covers both **Supervised Fine-Tuning (SFT)** and **Low-Rank Adaptation (LoRA)** methods for customizing pre-trained models. + +--- + +## πŸ“š Table of Contents + +- [πŸŽ“ Post-Training with Primus](#-post-training-with-primus) + - [πŸ“š Table of Contents](#-table-of-contents) + - [🎯 Overview](#-overview) + - [βš™οΈ Supported Backends](#️-supported-backends) + - [πŸ”§ Post-Training Methods](#-post-training-methods) + - [πŸš€ Quick Start](#-quick-start) + - [Prerequisites](#prerequisites) + - [Basic Usage](#basic-usage) + - [πŸ“ Configuration Examples](#-configuration-examples) + - [Supervised Fine-Tuning (SFT)](#supervised-fine-tuning-sft) + - [LoRA Fine-Tuning](#lora-fine-tuning) + - [πŸ–₯️ Single Node Training](#️-single-node-training) + - [Direct Mode](#direct-mode) + - [Container Mode](#container-mode) + - [πŸ“Š Hardware-Specific Configurations](#-hardware-specific-configurations) + - [MI300X Configurations](#mi300x-configurations) + - [MI355X Configurations](#mi355x-configurations) + - [🎨 Customizing Training Parameters](#-customizing-training-parameters) + - [πŸ’‘ Best Practices](#-best-practices) + - [πŸ” Troubleshooting](#-troubleshooting) + +--- + +## 🎯 Overview + +Post-training (fine-tuning) allows you to adapt pre-trained foundation models to specific tasks or domains. Primus supports two primary fine-tuning approaches: + +- **Supervised Fine-Tuning (SFT)**: Full fine-tuning that updates all model parameters +- **LoRA (Low-Rank Adaptation)**: Parameter-efficient fine-tuning that only trains lightweight adapter modules + +--- + +## βš™οΈ Supported Backends + +Post-training in Primus uses the **Megatron Bridge** backend: + +| Backend | Description | +| --------------- | --------------------------------------------------------------- | +| Megatron Bridge | Bridge implementation for fine-tuning Megatron-based models | + +--- + +## πŸ”§ Post-Training Methods + +| Method | Memory Usage | Training Speed | Use Case | +| ------ | ------------ | -------------- | ------------------------------------- | +| **SFT** | High | Slower | Maximum performance, full adaptation | +| **LoRA** | Low | Faster | Resource-efficient, quick iteration | + +**Key Differences:** +- **SFT** updates all model parameters, requiring more memory and compute +- **LoRA** trains only low-rank adapter matrices, significantly reducing resource requirements + +--- + +## πŸš€ Quick Start + +### Prerequisites + +- AMD ROCm drivers (β‰₯ 7.0) +- Docker (β‰₯ 24.0) with ROCm support (recommended) +- AMD Instinct GPUs (MI300X, MI355X, etc.) +- Pre-trained model checkpoint (optional, for continued training) + +```bash +# Quick verification +rocm-smi && docker --version +``` + +### Basic Usage + +The general command structure for post-training: + +```bash +./runner/primus-cli train posttrain --config +``` + +**Example commands:** + +```bash +# SFT with direct mode +./runner/primus-cli direct train posttrain \ + --config ./examples/megatron_bridge/configs/MI355X/qwen3_32b_sft_posttrain.yaml + +# LoRA with direct mode +./runner/primus-cli direct train posttrain \ + --config ./examples/megatron_bridge/configs/MI355X/qwen3_32b_lora_posttrain.yaml +``` + +--- + +## πŸ“ Configuration Examples + +### Supervised Fine-Tuning (SFT) + +Full fine-tuning configuration example for **Qwen3 32B** on **MI355X**: + +```yaml +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:qwen3_32b_sft_posttrain} +workspace: ${PRIMUS_WORKSPACE:./output} + +modules: + post_trainer: + framework: megatron_bridge + config: sft_trainer.yaml + model: qwen3_32b.yaml + + overrides: + # Parallelism configuration + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 1 + context_parallel_size: 1 + sequence_parallel: false + + # Fine-tuning method + peft: "none" # Full fine-tuning + + # Training configuration + train_iters: 200 + global_batch_size: 8 + micro_batch_size: 1 + seq_length: 8192 + + # Optimizer configuration + finetune_lr: 5.0e-6 + min_lr: 0.0 + lr_warmup_iters: 50 + + # Precision + precision_config: bf16_mixed +``` + +**Configuration location:** `examples/megatron_bridge/configs/MI355X/qwen3_32b_sft_posttrain.yaml` + +### LoRA Fine-Tuning + +Parameter-efficient fine-tuning configuration for **Qwen3 32B** on **MI355X**: + +```yaml +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:qwen3_32b_lora_posttrain} +workspace: ${PRIMUS_WORKSPACE:./output} + +modules: + post_trainer: + framework: megatron_bridge + config: sft_trainer.yaml + model: qwen3_32b.yaml + + overrides: + # Parallelism configuration + tensor_model_parallel_size: 1 # LoRA requires less parallelism + pipeline_model_parallel_size: 1 + context_parallel_size: 1 + sequence_parallel: false + + # Fine-tuning method + peft: lora # LoRA fine-tuning + + # Training configuration + train_iters: 200 + global_batch_size: 32 + micro_batch_size: 4 + seq_length: 8192 + + # Optimizer configuration + finetune_lr: 1.0e-4 # Higher LR for LoRA + min_lr: 0.0 + lr_warmup_iters: 50 + + # Precision + precision_config: bf16_mixed + + # Recompute configuration + recompute_granularity: full + recompute_method: uniform + recompute_num_layers: 1 +``` + +**Configuration location:** `examples/megatron_bridge/configs/MI355X/qwen3_32b_lora_posttrain.yaml` + +--- + +## πŸ–₯️ Single Node Training + +### Direct Mode + +Best for local development or when running directly on bare metal with ROCm installed. + +**SFT Example:** +```bash +./runner/primus-cli direct train posttrain \ + --config ./examples/megatron_bridge/configs/MI355X/qwen3_32b_sft_posttrain.yaml +``` + +**LoRA Example:** +```bash +./runner/primus-cli direct train posttrain \ + --config ./examples/megatron_bridge/configs/MI355X/qwen3_32b_lora_posttrain.yaml +``` + +**MI300X Examples:** +```bash +# SFT on MI300X +./runner/primus-cli direct train posttrain \ + --config ./examples/megatron_bridge/configs/MI300X/qwen3_32b_sft_posttrain.yaml + +# LoRA on MI300X +./runner/primus-cli direct train posttrain \ + --config ./examples/megatron_bridge/configs/MI300X/qwen3_32b_lora_posttrain.yaml +``` + +### Container Mode + +Recommended for environment isolation and dependency management. + +**Pull Docker image:** +```bash +docker pull docker.io/rocm/primus:latest +``` + +**SFT Example:** +```bash +./runner/primus-cli container --image rocm/primus:latest \ + train posttrain \ + --config ./examples/megatron_bridge/configs/MI355X/qwen3_32b_sft_posttrain.yaml +``` + +**LoRA Example:** +```bash +./runner/primus-cli container --image rocm/primus:latest \ + train posttrain \ + --config ./examples/megatron_bridge/configs/MI355X/qwen3_32b_lora_posttrain.yaml +``` + +--- + +## πŸ“Š Hardware-Specific Configurations + +### MI300X Configurations + +Available configurations for AMD Instinct MI300X GPUs: + +| Model | Method | Config File | TP | GBS | MBS | Seq Len | +| ---------- | ------ | ------------------------------------------- | -- | --- | --- | ------- | +| Qwen3 32B | SFT | `MI300X/qwen3_32b_sft_posttrain.yaml` | 2 | 8 | 2 | 8192 | +| Qwen3 32B | LoRA | `MI300X/qwen3_32b_lora_posttrain.yaml` | 1 | 32 | 2 | 8192 | + +**Example:** +```bash +./runner/primus-cli direct train posttrain \ + --config ./examples/megatron_bridge/configs/MI300X/qwen3_32b_sft_posttrain.yaml +``` + +### MI355X Configurations + +Available configurations for AMD Instinct MI355X GPUs: + +| Model | Method | Config File | TP | GBS | MBS | Seq Len | +| ------------ | ------ | ------------------------------------------- | -- | --- | --- | ------- | +| Qwen3 32B | SFT | `MI355X/qwen3_32b_sft_posttrain.yaml` | 1 | 8 | 1 | 8192 | +| Qwen3 32B | LoRA | `MI355X/qwen3_32b_lora_posttrain.yaml` | 1 | 32 | 4 | 8192 | + +**Legend:** +- **TP**: Tensor Parallelism Size +- **GBS**: Global Batch Size +- **MBS**: Micro Batch Size (per GPU) +- **Seq Len**: Sequence Length + +**Example:** +```bash +./runner/primus-cli direct train posttrain \ + --config ./examples/megatron_bridge/configs/MI355X/qwen3_32b_lora_posttrain.yaml +``` + +--- + +## 🎨 Customizing Training Parameters + +Key parameters you can customize in the YAML configuration: + +### Parallelism Settings +```yaml +tensor_model_parallel_size: 1 # Number of GPUs for tensor parallelism (1-8) +pipeline_model_parallel_size: 1 # Number of GPUs for pipeline parallelism +context_parallel_size: 1 # Context parallelism for long sequences +sequence_parallel: false # Enable sequence parallelism +``` + +### Training Hyperparameters +```yaml +train_iters: 200 # Total training iterations +global_batch_size: 8 # Global batch size (8-32 depending on config) +micro_batch_size: 1 # Batch size per GPU (1-4 depending on config) +seq_length: 2048 # Sequence length (2048-8192 depending on model) +eval_interval: 30 # Evaluate every N iterations +save_interval: 50 # Save checkpoint every N iterations +``` + +### Learning Rate Configuration +```yaml +finetune_lr: 1.0e-4 # Initial learning rate +min_lr: 0.0 # Minimum learning rate +lr_warmup_iters: 50 # Number of warmup iterations +lr_decay_iters: null # Learning rate decay iterations +``` + +### Fine-Tuning Method +```yaml +peft: lora # Options: "lora" or "none" (for full SFT) +packed_sequence: false # Enable packed sequences for efficiency +``` + +### Precision Configuration +```yaml +precision_config: bf16_mixed # Options: bf16_mixed, fp16_mixed, fp32 +``` + +### Memory Optimization +```yaml +recompute_granularity: full # Options: full, selective, null +recompute_method: uniform # Recompute strategy +recompute_num_layers: 1 # Number of layers to recompute +``` + +--- + +## πŸ’‘ Best Practices + +### Choosing Between SFT and LoRA + +**Use SFT when:** +- You need maximum model performance +- You have sufficient GPU memory +- Training time is not critical +- You want full model adaptation + +**Use LoRA when:** +- GPU memory is limited +- You need fast iteration cycles +- Training multiple task-specific adapters +- Parameter efficiency is important + +### Parallelism Configuration + +**For SFT:** +- Use higher `tensor_model_parallel_size` for large models (e.g., TP=8 for 70B) +- Consider pipeline parallelism for very large models +- Examples: + - 32B model: TP=1-2 (MI300X: TP=2, MI355X: TP=1) + - 70B model: TP=8 + +**For LoRA:** +- Lower `tensor_model_parallel_size` due to reduced memory +- LoRA can fit larger models with less parallelism +- Examples: + - 32B model: TP=1 + - 70B model: TP=8 (still requires high TP due to model size) + +### Learning Rate Guidelines + +- **SFT**: Use lower learning rates (5e-6 to 1e-5) +- **LoRA**: Use higher learning rates (1e-4 to 5e-4) +- Always use warmup for stable training + +### Batch Size Recommendations + +- Start with `global_batch_size: 8` for SFT development +- LoRA can use higher batch sizes (e.g., 32) due to lower memory usage +- Increase for production: 64, 128, or higher +- Adjust `micro_batch_size` (1-4) based on GPU memory and sequence length +- Longer sequences (8192) may require higher `micro_batch_size` for efficiency + +--- + +## πŸ” Troubleshooting + +### Out of Memory (OOM) Errors + +**For SFT:** +1. Increase `tensor_model_parallel_size` +2. Reduce `micro_batch_size` +3. Enable gradient checkpointing: + ```yaml + recompute_granularity: full + recompute_method: uniform + recompute_num_layers: 1 + ``` +4. Reduce `seq_length` + +**For LoRA:** +1. LoRA should have lower memory usage; verify `peft: lora` is set +2. Reduce `micro_batch_size` if still facing OOM +3. Enable recomputation as above + +### Training Instability + +1. **Check learning rate**: Reduce if loss is spiking +2. **Increase warmup**: Try `lr_warmup_iters: 100` or higher +3. **Use mixed precision**: Ensure `precision_config: bf16_mixed` +4. **Monitor gradients**: Watch for gradient explosions + +### Slow Training Speed + +1. **Optimize batch size**: Increase `global_batch_size` if possible +2. **Check parallelism**: Ensure optimal TP/PP configuration +3. **Use container mode**: Docker containers can improve performance +4. **Profile execution**: Use profiling tools to identify bottlenecks + +### Configuration Issues + +1. **Verify paths**: Ensure config file paths are correct +2. **Check YAML syntax**: Validate indentation and structure +3. **Environment variables**: Set `PRIMUS_WORKSPACE` if needed +4. **Model checkpoint**: Verify pre-trained checkpoint path (if using) + +--- + +## 🎯 Summary Commands + +**Quick reference for common post-training tasks:** + +```bash +# SFT on MI355X (direct mode) +./runner/primus-cli direct train posttrain \ + --config ./examples/megatron_bridge/configs/MI355X/qwen3_32b_sft_posttrain.yaml + +# LoRA on MI355X (direct mode) +./runner/primus-cli direct train posttrain \ + --config ./examples/megatron_bridge/configs/MI355X/qwen3_32b_lora_posttrain.yaml + +# SFT on MI300X (container mode) +./runner/primus-cli container --image rocm/primus:latest train posttrain \ + --config ./examples/megatron_bridge/configs/MI300X/qwen3_32b_sft_posttrain.yaml +``` + +--- + +**Need help?** Open an issue on [GitHub](https://github.com/AMD-AIG-AIMA/Primus/issues). + +**Start fine-tuning with Primus! πŸš€** diff --git a/examples/megatron_bridge/configs/MI300X/qwen3_32b_lora_posttrain.yaml b/examples/megatron_bridge/configs/MI300X/qwen3_32b_lora_posttrain.yaml new file mode 100644 index 000000000..456ee968a --- /dev/null +++ b/examples/megatron_bridge/configs/MI300X/qwen3_32b_lora_posttrain.yaml @@ -0,0 +1,58 @@ +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:qwen3_32b_lora_posttrain} +workspace: ${PRIMUS_WORKSPACE:./output} + +modules: + post_trainer: + framework: megatron_bridge + config: sft_trainer.yaml + + # Model to run + model: qwen3_32b.yaml + + overrides: + stderr_sink_level: DEBUG + + # Parallelism configuration + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 1 + pipeline_dtype: null + virtual_pipeline_model_parallel_size: null + context_parallel_size: 1 + sequence_parallel: false + use_megatron_fsdp: false + + # Finetuning-specific params + #pretrained_checkpoint: null + peft: lora + packed_sequence: false + + # Training configuration + train_iters: 200 + global_batch_size: 32 + micro_batch_size: 2 + seq_length: 8192 + eval_interval: 30 + save_interval: 50 + + # Optimizer configuration + finetune_lr: 1.0e-4 + min_lr: 0.0 + lr_warmup_iters: 50 + lr_decay_iters: null + + # W&B logging + wandb_project: null + wandb_entity: null + wandb_exp_name: null + + # Precision + precision_config: bf16_mixed + comm_overlap_config: null + + # Recompute configuration (enabled for 32B model) + recompute_granularity: full + recompute_method: uniform + recompute_num_layers: 1 + diff --git a/examples/megatron_bridge/configs/MI300X/qwen3_32b_sft_posttrain.yaml b/examples/megatron_bridge/configs/MI300X/qwen3_32b_sft_posttrain.yaml new file mode 100644 index 000000000..c91a6b0d5 --- /dev/null +++ b/examples/megatron_bridge/configs/MI300X/qwen3_32b_sft_posttrain.yaml @@ -0,0 +1,58 @@ +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:qwen3_32b_sft_posttrain} +workspace: ${PRIMUS_WORKSPACE:./output} + +modules: + post_trainer: + framework: megatron_bridge + config: sft_trainer.yaml + + # Model to run + model: qwen3_32b.yaml + + overrides: + stderr_sink_level: DEBUG + + # Parallelism configuration + tensor_model_parallel_size: 2 + pipeline_model_parallel_size: 1 + pipeline_dtype: null + virtual_pipeline_model_parallel_size: null + context_parallel_size: 1 + sequence_parallel: false + use_megatron_fsdp: false + + # Finetuning-specific params + #pretrained_checkpoint: null + peft: "none" + packed_sequence: false + + # Training configuration + train_iters: 200 + global_batch_size: 8 + micro_batch_size: 2 + seq_length: 8192 + eval_interval: 30 + save_interval: 50 + + # Optimizer configuration + finetune_lr: 5.0e-6 + min_lr: 0.0 + lr_warmup_iters: 50 + lr_decay_iters: null + + # W&B logging + wandb_project: null + wandb_entity: null + wandb_exp_name: null + + # Precision + precision_config: bf16_mixed + comm_overlap_config: null + + # Recompute configuration (enabled for 32B model) + recompute_granularity: full + recompute_method: uniform + recompute_num_layers: 1 + diff --git a/examples/megatron_bridge/configs/MI355X/qwen3_32b_lora_posttrain.yaml b/examples/megatron_bridge/configs/MI355X/qwen3_32b_lora_posttrain.yaml index 49ce4f8a0..3dfb3eb39 100755 --- a/examples/megatron_bridge/configs/MI355X/qwen3_32b_lora_posttrain.yaml +++ b/examples/megatron_bridge/configs/MI355X/qwen3_32b_lora_posttrain.yaml @@ -30,9 +30,9 @@ modules: # Training configuration train_iters: 200 - global_batch_size: 128 - micro_batch_size: 1 - seq_length: 2048 + global_batch_size: 32 + micro_batch_size: 4 + seq_length: 8192 eval_interval: 30 save_interval: 50 diff --git a/examples/megatron_bridge/configs/MI355X/qwen3_32b_sft_posttrain.yaml b/examples/megatron_bridge/configs/MI355X/qwen3_32b_sft_posttrain.yaml index 1b9f10aff..d01623990 100755 --- a/examples/megatron_bridge/configs/MI355X/qwen3_32b_sft_posttrain.yaml +++ b/examples/megatron_bridge/configs/MI355X/qwen3_32b_sft_posttrain.yaml @@ -15,7 +15,7 @@ modules: stderr_sink_level: DEBUG # Parallelism configuration - tensor_model_parallel_size: 4 + tensor_model_parallel_size: 1 pipeline_model_parallel_size: 1 pipeline_dtype: null virtual_pipeline_model_parallel_size: null @@ -30,9 +30,9 @@ modules: # Training configuration train_iters: 200 - global_batch_size: 128 + global_batch_size: 8 micro_batch_size: 1 - seq_length: 2048 + seq_length: 8192 eval_interval: 30 save_interval: 50 diff --git a/primus/configs/modules/megatron_bridge/sft_trainer.yaml b/primus/configs/modules/megatron_bridge/sft_trainer.yaml index 230bb22a5..994a1389a 100644 --- a/primus/configs/modules/megatron_bridge/sft_trainer.yaml +++ b/primus/configs/modules/megatron_bridge/sft_trainer.yaml @@ -12,7 +12,7 @@ stage: "sft" # main control flag -enable_primus_turbo: false +enable_primus_turbo: true # feature control flags use_turbo_attention: false From 8db36dced200d4ba1cd519713b4e059961ed0d71 Mon Sep 17 00:00:00 2001 From: vidushi8 Date: Wed, 25 Feb 2026 18:47:20 -0600 Subject: [PATCH 03/11] update MI355 yaml for better perf --- examples/megatron/configs/MI355X/llama2_7B-FP8-pretrain.yaml | 2 +- .../megatron/configs/MI355X/llama3.1_70B-FP8-pretrain.yaml | 4 ++-- .../megatron/configs/MI355X/llama3.1_8B-FP8-pretrain.yaml | 2 +- examples/megatron/configs/MI355X/qwen2.5_7B-FP8-pretrain.yaml | 2 +- .../torchtitan/configs/MI355X/llama3.1_70B-FP8-pretrain.yaml | 1 + 5 files changed, 6 insertions(+), 5 deletions(-) diff --git a/examples/megatron/configs/MI355X/llama2_7B-FP8-pretrain.yaml b/examples/megatron/configs/MI355X/llama2_7B-FP8-pretrain.yaml index b77adb414..3c390c575 100755 --- a/examples/megatron/configs/MI355X/llama2_7B-FP8-pretrain.yaml +++ b/examples/megatron/configs/MI355X/llama2_7B-FP8-pretrain.yaml @@ -47,7 +47,7 @@ modules: expert_model_parallel_size: 1 overlap_grad_reduce: true overlap_param_gather: true - gradient_accumulation_fusion: false + gradient_accumulation_fusion: true # data mock_data: true diff --git a/examples/megatron/configs/MI355X/llama3.1_70B-FP8-pretrain.yaml b/examples/megatron/configs/MI355X/llama3.1_70B-FP8-pretrain.yaml index 2c2fa0b68..9975256f7 100644 --- a/examples/megatron/configs/MI355X/llama3.1_70B-FP8-pretrain.yaml +++ b/examples/megatron/configs/MI355X/llama3.1_70B-FP8-pretrain.yaml @@ -17,8 +17,8 @@ modules: log_avg_reset_interval: 50 train_iters: 50 - micro_batch_size: 3 - global_batch_size: 24 + micro_batch_size: 4 + global_batch_size: 32 seq_length: 8192 max_position_embeddings: 8192 diff --git a/examples/megatron/configs/MI355X/llama3.1_8B-FP8-pretrain.yaml b/examples/megatron/configs/MI355X/llama3.1_8B-FP8-pretrain.yaml index cff682659..db6b7d0b1 100644 --- a/examples/megatron/configs/MI355X/llama3.1_8B-FP8-pretrain.yaml +++ b/examples/megatron/configs/MI355X/llama3.1_8B-FP8-pretrain.yaml @@ -45,7 +45,7 @@ modules: expert_model_parallel_size: 1 overlap_grad_reduce: true overlap_param_gather: true - gradient_accumulation_fusion: false + gradient_accumulation_fusion: true # data mock_data: true diff --git a/examples/megatron/configs/MI355X/qwen2.5_7B-FP8-pretrain.yaml b/examples/megatron/configs/MI355X/qwen2.5_7B-FP8-pretrain.yaml index 872180b34..55c0762f9 100644 --- a/examples/megatron/configs/MI355X/qwen2.5_7B-FP8-pretrain.yaml +++ b/examples/megatron/configs/MI355X/qwen2.5_7B-FP8-pretrain.yaml @@ -46,7 +46,7 @@ modules: expert_model_parallel_size: 1 overlap_grad_reduce: true overlap_param_gather: true - gradient_accumulation_fusion: false + gradient_accumulation_fusion: true # data mock_data: true diff --git a/examples/torchtitan/configs/MI355X/llama3.1_70B-FP8-pretrain.yaml b/examples/torchtitan/configs/MI355X/llama3.1_70B-FP8-pretrain.yaml index 91b92ed32..954b3826d 100644 --- a/examples/torchtitan/configs/MI355X/llama3.1_70B-FP8-pretrain.yaml +++ b/examples/torchtitan/configs/MI355X/llama3.1_70B-FP8-pretrain.yaml @@ -29,6 +29,7 @@ modules: training: local_batch_size: 3 + global_batch_size: 96 seq_len: 8192 mock_data: false steps: 10 From 143d59305801c12aaffe77e32cb965fb23be3a11 Mon Sep 17 00:00:00 2001 From: vidushi8 Date: Thu, 26 Feb 2026 15:41:09 -0600 Subject: [PATCH 04/11] update yaml --- .../configs/MI355X/deepseek_v2_lite-BF16-pretrain.yaml | 3 ++- .../megatron/configs/MI355X/llama3.1_70B-BF16-pretrain.yaml | 4 ++-- .../megatron/configs/MI355X/llama3.1_8B-FP8-pretrain.yaml | 4 ++-- 3 files changed, 6 insertions(+), 5 deletions(-) diff --git a/examples/megatron/configs/MI355X/deepseek_v2_lite-BF16-pretrain.yaml b/examples/megatron/configs/MI355X/deepseek_v2_lite-BF16-pretrain.yaml index cf8064671..54051c9c7 100644 --- a/examples/megatron/configs/MI355X/deepseek_v2_lite-BF16-pretrain.yaml +++ b/examples/megatron/configs/MI355X/deepseek_v2_lite-BF16-pretrain.yaml @@ -73,7 +73,7 @@ modules: # Turbo enable_primus_turbo: true use_turbo_attention: false - use_turbo_grouped_mlp: false + use_turbo_grouped_mlp: true # deepep use_turbo_deepep: true @@ -87,6 +87,7 @@ modules: # sync-free moe support stage 1-2, 0 means not use sync-free moe # stage 2 is recommended for better performance turbo_sync_free_moe_stage: 1 + moe_use_fused_router_with_aux_score: true # Cross entropy flags cross_entropy_fusion_impl: "te" diff --git a/examples/megatron/configs/MI355X/llama3.1_70B-BF16-pretrain.yaml b/examples/megatron/configs/MI355X/llama3.1_70B-BF16-pretrain.yaml index 6ee1f25b3..c707572da 100644 --- a/examples/megatron/configs/MI355X/llama3.1_70B-BF16-pretrain.yaml +++ b/examples/megatron/configs/MI355X/llama3.1_70B-BF16-pretrain.yaml @@ -17,8 +17,8 @@ modules: log_avg_reset_interval: 50 train_iters: 50 - micro_batch_size: 4 - global_batch_size: 32 + micro_batch_size: 7 + global_batch_size: 56 seq_length: 8192 max_position_embeddings: 8192 diff --git a/examples/megatron/configs/MI355X/llama3.1_8B-FP8-pretrain.yaml b/examples/megatron/configs/MI355X/llama3.1_8B-FP8-pretrain.yaml index db6b7d0b1..1d8dcd9fc 100644 --- a/examples/megatron/configs/MI355X/llama3.1_8B-FP8-pretrain.yaml +++ b/examples/megatron/configs/MI355X/llama3.1_8B-FP8-pretrain.yaml @@ -21,8 +21,8 @@ modules: log_avg_reset_interval: 50 train_iters: 50 - micro_batch_size: 4 - global_batch_size: 256 + micro_batch_size: 6 + global_batch_size: 384 seq_length: 8192 max_position_embeddings: 8192 From 3dade7140303116324d9f1aebdc22a4cc2705f5d Mon Sep 17 00:00:00 2001 From: clairesonglee Date: Thu, 26 Feb 2026 13:57:58 -0800 Subject: [PATCH 05/11] [Primus] Fix MoE MLA issue from hybrid models branch merge (#564) Co-authored-by: Mingyu Yang Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com> Co-authored-by: Kailash Gogineni Co-authored-by: HuangWei-95 Co-authored-by: HuangWei-95 Co-authored-by: Xiaoming-AMD Co-authored-by: WangLingxun --- .../configs/MI300X/mamba_370M-pretrain.yaml | 85 ++++ .../MI300X/zebra_llama_1B-pretrain.yaml | 70 +++ .../MI300X/zebra_llama_3B-pretrain.yaml | 70 +++ .../MI300X/zebra_llama_8B-pretrain.yaml | 70 +++ .../MI355X/zebra_llama_1B-pretrain.yaml | 70 +++ .../MI355X/zebra_llama_3B-pretrain.yaml | 70 +++ .../MI355X/zebra_llama_8B-pretrain.yaml | 70 +++ .../megatron/core/models/hybrid/__init__.py | 10 + .../core/models/hybrid/hybrid_block.py | 403 ++++++++++++++++++ .../hybrid/hybrid_mamba_mla_layer_specs.py | 119 ++++++ .../megatron/megatron_pretrain_trainer.py | 35 +- .../models/megatron/deepseek_v2_base.yaml | 1 + .../models/megatron/deepseek_v3_base.yaml | 1 + .../models/megatron/hybrid_model_base.yaml | 21 + .../models/megatron/language_model.yaml | 6 +- .../configs/models/megatron/mamba_370M.yaml | 16 + .../configs/models/megatron/mamba_base.yaml | 36 ++ .../models/megatron/zebra_llama_1B.yaml | 42 ++ .../models/megatron/zebra_llama_3B.yaml | 43 ++ .../models/megatron/zebra_llama_8B.yaml | 43 ++ .../modules/megatron/trainer_base.yaml | 23 +- primus/core/utils/import_utils.py | 43 +- .../modules/trainer/megatron/pre_trainer.py | 33 +- primus/modules/trainer/megatron/trainer.py | 36 +- 24 files changed, 1375 insertions(+), 41 deletions(-) create mode 100644 examples/megatron/configs/MI300X/mamba_370M-pretrain.yaml create mode 100644 examples/megatron/configs/MI300X/zebra_llama_1B-pretrain.yaml create mode 100644 examples/megatron/configs/MI300X/zebra_llama_3B-pretrain.yaml create mode 100644 examples/megatron/configs/MI300X/zebra_llama_8B-pretrain.yaml create mode 100644 examples/megatron/configs/MI355X/zebra_llama_1B-pretrain.yaml create mode 100644 examples/megatron/configs/MI355X/zebra_llama_3B-pretrain.yaml create mode 100644 examples/megatron/configs/MI355X/zebra_llama_8B-pretrain.yaml create mode 100644 primus/backends/megatron/core/models/hybrid/__init__.py create mode 100644 primus/backends/megatron/core/models/hybrid/hybrid_block.py create mode 100644 primus/backends/megatron/core/models/hybrid/hybrid_mamba_mla_layer_specs.py create mode 100644 primus/configs/models/megatron/hybrid_model_base.yaml create mode 100644 primus/configs/models/megatron/mamba_370M.yaml create mode 100644 primus/configs/models/megatron/mamba_base.yaml create mode 100644 primus/configs/models/megatron/zebra_llama_1B.yaml create mode 100644 primus/configs/models/megatron/zebra_llama_3B.yaml create mode 100644 primus/configs/models/megatron/zebra_llama_8B.yaml diff --git a/examples/megatron/configs/MI300X/mamba_370M-pretrain.yaml b/examples/megatron/configs/MI300X/mamba_370M-pretrain.yaml new file mode 100644 index 000000000..469913761 --- /dev/null +++ b/examples/megatron/configs/MI300X/mamba_370M-pretrain.yaml @@ -0,0 +1,85 @@ +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:mamba_370M-pretrain} +workspace: ${PRIMUS_WORKSPACE:./output} + +modules: + pre_trainer: + framework: megatron + config: pre_trainer.yaml + + # model to run + model: mamba_370M.yaml + overrides: + # log + wandb_project: "Primus_Mamba_Pretrain" + # disable_wandb: false + # disable_tensorboard: false + stderr_sink_level: DEBUG + + eval_iters: 0 + + log_avg_skip_iterations: 2 + log_avg_reset_interval: 50 + + train_iters: 50 + micro_batch_size: 4 + global_batch_size: 256 + + seq_length: 2048 + max_position_embeddings: 2048 + + lr: 3.0e-4 + min_lr: 0.0 + lr_warmup_iters: 50000 + lr_decay_iters: 73192188 + lr_decay_style: cosine + weight_decay: 0.1 + adam_beta1: 0.9 + adam_beta2: 0.95 + eod_mask_loss: true + init_method_std: 0.02 + norm_epsilon: 1.0e-5 + + # Mamba-specific: must provide spec + spec: ['megatron.core.models.mamba.mamba_layer_specs', 'mamba_stack_spec'] + + # Tokenizer + tokenizer_type: HuggingFaceTokenizer + tokenizer_model: EleutherAI/gpt-neox-20b + + # parallel + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 1 + expert_model_parallel_size: 1 + overlap_grad_reduce: true + overlap_param_gather: true + gradient_accumulation_fusion: false + + # data + mock_data: true + train_data_path: null + valid_data_path: null + test_data_path: null + + # ckpt + finetune: false + auto_continue_train: false + load: null + no_load_optim: null + no_load_rng: null + save: null + save_interval: 20000 + no_save_optim: null + no_save_rng: null + disable_last_saving: true + ckpt_format: torch + + # Turbo - may need to disable for Mamba if not supported + enable_primus_turbo: false + use_turbo_attention: false + use_turbo_grouped_mlp: false + + # Cross entropy flags + # cross_entropy_fusion_impl: "native" + # cross_entropy_loss_fusion: false diff --git a/examples/megatron/configs/MI300X/zebra_llama_1B-pretrain.yaml b/examples/megatron/configs/MI300X/zebra_llama_1B-pretrain.yaml new file mode 100644 index 000000000..d2327bab8 --- /dev/null +++ b/examples/megatron/configs/MI300X/zebra_llama_1B-pretrain.yaml @@ -0,0 +1,70 @@ +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:zebra_llama_1B-pretrain} +workspace: ${PRIMUS_WORKSPACE:./output} + +modules: + pre_trainer: + framework: megatron + config: pre_trainer.yaml + + # model to run + model: zebra_llama_1B.yaml + overrides: + # log + wandb_project: "Primus_Zebra_Llama_1B_Pretrain" + stderr_sink_level: DEBUG + + eval_iters: 0 + + log_avg_skip_iterations: 2 + log_avg_reset_interval: 50 + + train_iters: 100 + micro_batch_size: 2 + global_batch_size: 16 + + seq_length: 8192 + max_position_embeddings: 8192 + original_max_position_embeddings: 8192 + + lr: 2.0e-4 + min_lr: 2.0e-5 + lr_warmup_iters: 200 + lr_decay_iters: 10000 + lr_decay_style: cosine + weight_decay: 0.1 + adam_beta1: 0.9 + adam_beta2: 0.95 + eod_mask_loss: true + + # Mamba-specific: must provide spec + # Use custom hybrid Mamba+MLA spec + spec: ['primus.backends.megatron.core.models.hybrid.hybrid_mamba_mla_layer_specs', 'hybrid_stack_spec'] + + # Tokenizer + tokenizer_type: HuggingFaceTokenizer + tokenizer_model: meta-llama/Llama-3.2-1B + + # parallel + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 1 + expert_model_parallel_size: 1 + overlap_grad_reduce: true + overlap_param_gather: true + gradient_accumulation_fusion: false + + # data + mock_data: true + train_data_path: null + valid_data_path: null + test_data_path: null + + # ckpt + finetune: false + auto_continue_train: false + load: null + save: null + save_interval: 10000 + disable_last_saving: true + ckpt_format: torch diff --git a/examples/megatron/configs/MI300X/zebra_llama_3B-pretrain.yaml b/examples/megatron/configs/MI300X/zebra_llama_3B-pretrain.yaml new file mode 100644 index 000000000..4daeb25e6 --- /dev/null +++ b/examples/megatron/configs/MI300X/zebra_llama_3B-pretrain.yaml @@ -0,0 +1,70 @@ +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:zebra_llama_3B-pretrain} +workspace: ${PRIMUS_WORKSPACE:./output} + +modules: + pre_trainer: + framework: megatron + config: pre_trainer.yaml + + # model to run + model: zebra_llama_3B.yaml + overrides: + # log + wandb_project: "Primus_Zebra_Llama_3B_Pretrain" + stderr_sink_level: DEBUG + + eval_iters: 0 + + log_avg_skip_iterations: 2 + log_avg_reset_interval: 50 + + train_iters: 100 + micro_batch_size: 2 + global_batch_size: 16 + + seq_length: 8192 + max_position_embeddings: 8192 + original_max_position_embeddings: 8192 + + lr: 2.0e-4 + min_lr: 2.0e-5 + lr_warmup_iters: 200 + lr_decay_iters: 10000 + lr_decay_style: cosine + weight_decay: 0.1 + adam_beta1: 0.9 + adam_beta2: 0.95 + eod_mask_loss: true + + # Mamba-specific: must provide spec + # Use custom hybrid Mamba+MLA spec + spec: ['primus.backends.megatron.core.models.hybrid.hybrid_mamba_mla_layer_specs', 'hybrid_stack_spec'] + + # Tokenizer + tokenizer_type: HuggingFaceTokenizer + tokenizer_model: meta-llama/Llama-3.2-3B + + # parallel + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 1 + expert_model_parallel_size: 1 + overlap_grad_reduce: true + overlap_param_gather: true + gradient_accumulation_fusion: false + + # data + mock_data: true + train_data_path: null + valid_data_path: null + test_data_path: null + + # ckpt + finetune: false + auto_continue_train: false + load: null + save: null + save_interval: 10000 + disable_last_saving: true + ckpt_format: torch diff --git a/examples/megatron/configs/MI300X/zebra_llama_8B-pretrain.yaml b/examples/megatron/configs/MI300X/zebra_llama_8B-pretrain.yaml new file mode 100644 index 000000000..a7083c069 --- /dev/null +++ b/examples/megatron/configs/MI300X/zebra_llama_8B-pretrain.yaml @@ -0,0 +1,70 @@ +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:zebra_llama_8B-pretrain} +workspace: ${PRIMUS_WORKSPACE:./output} + +modules: + pre_trainer: + framework: megatron + config: pre_trainer.yaml + + # model to run + model: zebra_llama_8B.yaml + overrides: + # log + wandb_project: "Primus_Zebra_Llama_8B_Pretrain" + stderr_sink_level: DEBUG + + eval_iters: 0 + + log_avg_skip_iterations: 2 + log_avg_reset_interval: 50 + + train_iters: 100 + micro_batch_size: 2 + global_batch_size: 16 + + seq_length: 8192 + max_position_embeddings: 8192 + original_max_position_embeddings: 8192 + + lr: 2.0e-4 + min_lr: 2.0e-5 + lr_warmup_iters: 200 + lr_decay_iters: 10000 + lr_decay_style: cosine + weight_decay: 0.1 + adam_beta1: 0.9 + adam_beta2: 0.95 + eod_mask_loss: true + + # Mamba-specific: must provide spec + # Use custom hybrid Mamba+MLA spec + spec: ['primus.backends.megatron.core.models.hybrid.hybrid_mamba_mla_layer_specs', 'hybrid_stack_spec'] + + # Tokenizer + tokenizer_type: HuggingFaceTokenizer + tokenizer_model: meta-llama/Llama-3.1-8B + + # parallel + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 1 + expert_model_parallel_size: 1 + overlap_grad_reduce: true + overlap_param_gather: true + gradient_accumulation_fusion: false + + # data + mock_data: true + train_data_path: null + valid_data_path: null + test_data_path: null + + # ckpt + finetune: false + auto_continue_train: false + load: null + save: null + save_interval: 10000 + disable_last_saving: true + ckpt_format: torch diff --git a/examples/megatron/configs/MI355X/zebra_llama_1B-pretrain.yaml b/examples/megatron/configs/MI355X/zebra_llama_1B-pretrain.yaml new file mode 100644 index 000000000..4eaf5102c --- /dev/null +++ b/examples/megatron/configs/MI355X/zebra_llama_1B-pretrain.yaml @@ -0,0 +1,70 @@ +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:zebra_llama_1B-pretrain} +workspace: ${PRIMUS_WORKSPACE:./output} + +modules: + pre_trainer: + framework: megatron + config: pre_trainer.yaml + + # model to run + model: zebra_llama_1B.yaml + overrides: + # log + wandb_project: "Primus_Zebra_Llama_1B_Pretrain" + stderr_sink_level: DEBUG + + eval_iters: 0 + + log_avg_skip_iterations: 2 + log_avg_reset_interval: 50 + + train_iters: 100 + micro_batch_size: 16 + global_batch_size: 128 + + seq_length: 8192 + max_position_embeddings: 8192 + original_max_position_embeddings: 8192 + + lr: 2.0e-4 + min_lr: 2.0e-5 + lr_warmup_iters: 200 + lr_decay_iters: 10000 + lr_decay_style: cosine + weight_decay: 0.1 + adam_beta1: 0.9 + adam_beta2: 0.95 + eod_mask_loss: true + + # Mamba-specific: must provide spec + # Use custom hybrid Mamba+MLA spec + spec: ['primus.backends.megatron.core.models.hybrid.hybrid_mamba_mla_layer_specs', 'hybrid_stack_spec'] + + # Tokenizer + tokenizer_type: HuggingFaceTokenizer + tokenizer_model: meta-llama/Llama-3.2-1B + + # parallel + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 1 + expert_model_parallel_size: 1 + overlap_grad_reduce: true + overlap_param_gather: true + gradient_accumulation_fusion: false + + # data + mock_data: true + train_data_path: null + valid_data_path: null + test_data_path: null + + # ckpt + finetune: false + auto_continue_train: false + load: null + save: null + save_interval: 10000 + disable_last_saving: true + ckpt_format: torch diff --git a/examples/megatron/configs/MI355X/zebra_llama_3B-pretrain.yaml b/examples/megatron/configs/MI355X/zebra_llama_3B-pretrain.yaml new file mode 100644 index 000000000..4daeb25e6 --- /dev/null +++ b/examples/megatron/configs/MI355X/zebra_llama_3B-pretrain.yaml @@ -0,0 +1,70 @@ +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:zebra_llama_3B-pretrain} +workspace: ${PRIMUS_WORKSPACE:./output} + +modules: + pre_trainer: + framework: megatron + config: pre_trainer.yaml + + # model to run + model: zebra_llama_3B.yaml + overrides: + # log + wandb_project: "Primus_Zebra_Llama_3B_Pretrain" + stderr_sink_level: DEBUG + + eval_iters: 0 + + log_avg_skip_iterations: 2 + log_avg_reset_interval: 50 + + train_iters: 100 + micro_batch_size: 2 + global_batch_size: 16 + + seq_length: 8192 + max_position_embeddings: 8192 + original_max_position_embeddings: 8192 + + lr: 2.0e-4 + min_lr: 2.0e-5 + lr_warmup_iters: 200 + lr_decay_iters: 10000 + lr_decay_style: cosine + weight_decay: 0.1 + adam_beta1: 0.9 + adam_beta2: 0.95 + eod_mask_loss: true + + # Mamba-specific: must provide spec + # Use custom hybrid Mamba+MLA spec + spec: ['primus.backends.megatron.core.models.hybrid.hybrid_mamba_mla_layer_specs', 'hybrid_stack_spec'] + + # Tokenizer + tokenizer_type: HuggingFaceTokenizer + tokenizer_model: meta-llama/Llama-3.2-3B + + # parallel + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 1 + expert_model_parallel_size: 1 + overlap_grad_reduce: true + overlap_param_gather: true + gradient_accumulation_fusion: false + + # data + mock_data: true + train_data_path: null + valid_data_path: null + test_data_path: null + + # ckpt + finetune: false + auto_continue_train: false + load: null + save: null + save_interval: 10000 + disable_last_saving: true + ckpt_format: torch diff --git a/examples/megatron/configs/MI355X/zebra_llama_8B-pretrain.yaml b/examples/megatron/configs/MI355X/zebra_llama_8B-pretrain.yaml new file mode 100644 index 000000000..ff274291e --- /dev/null +++ b/examples/megatron/configs/MI355X/zebra_llama_8B-pretrain.yaml @@ -0,0 +1,70 @@ +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:zebra_llama_8B-pretrain} +workspace: ${PRIMUS_WORKSPACE:./output} + +modules: + pre_trainer: + framework: megatron + config: pre_trainer.yaml + + # model to run + model: zebra_llama_8B.yaml + overrides: + # log + wandb_project: "Primus_Zebra_Llama_8B_Pretrain" + stderr_sink_level: DEBUG + + eval_iters: 0 + + log_avg_skip_iterations: 2 + log_avg_reset_interval: 50 + + train_iters: 100 + micro_batch_size: 16 + global_batch_size: 128 + + seq_length: 8192 + max_position_embeddings: 8192 + original_max_position_embeddings: 8192 + + lr: 2.0e-4 + min_lr: 2.0e-5 + lr_warmup_iters: 200 + lr_decay_iters: 10000 + lr_decay_style: cosine + weight_decay: 0.1 + adam_beta1: 0.9 + adam_beta2: 0.95 + eod_mask_loss: true + + # Mamba-specific: must provide spec + # Use custom hybrid Mamba+MLA spec + spec: ['primus.backends.megatron.core.models.hybrid.hybrid_mamba_mla_layer_specs', 'hybrid_stack_spec'] + + # Tokenizer + tokenizer_type: HuggingFaceTokenizer + tokenizer_model: meta-llama/Llama-3.1-8B + + # parallel + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 1 + expert_model_parallel_size: 1 + overlap_grad_reduce: true + overlap_param_gather: true + gradient_accumulation_fusion: false + + # data + mock_data: true + train_data_path: null + valid_data_path: null + test_data_path: null + + # ckpt + finetune: false + auto_continue_train: false + load: null + save: null + save_interval: 10000 + disable_last_saving: true + ckpt_format: torch diff --git a/primus/backends/megatron/core/models/hybrid/__init__.py b/primus/backends/megatron/core/models/hybrid/__init__.py new file mode 100644 index 000000000..05d2d673a --- /dev/null +++ b/primus/backends/megatron/core/models/hybrid/__init__.py @@ -0,0 +1,10 @@ +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Hybrid Mamba+MLA layer specifications for Megatron-LM.""" + +from .hybrid_mamba_mla_layer_specs import hybrid_stack_spec + +__all__ = ["hybrid_stack_spec"] diff --git a/primus/backends/megatron/core/models/hybrid/hybrid_block.py b/primus/backends/megatron/core/models/hybrid/hybrid_block.py new file mode 100644 index 000000000..3c4d7def8 --- /dev/null +++ b/primus/backends/megatron/core/models/hybrid/hybrid_block.py @@ -0,0 +1,403 @@ +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2024, Tri Dao, Albert Gu. + +# Some of this code was adopted from https://github.com/state-spaces/mamba/ +# This source code is licensed under the Apache license found in the +# LICENSE file in the root directory of this source tree. + +from contextlib import nullcontext +from dataclasses import dataclass +from typing import Optional, Tuple, Union + +import torch +from megatron.core.dist_checkpointing.mapping import ShardedStateDict +from megatron.core.dist_checkpointing.utils import replace_prefix_for_sharding +from megatron.core.enums import Fp8Recipe +from megatron.core.extensions.transformer_engine import TENorm +from megatron.core.fp8_utils import get_fp8_context +from megatron.core.inference.contexts import BaseInferenceContext +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.ssm.mamba_hybrid_layer_allocation import Symbols as LayerSymbols +from megatron.core.transformer import TransformerConfig +from torch import Tensor, nn + +# CudaGraphScope is not available in older Megatron versions +try: + from megatron.core.transformer.enums import CudaGraphScope + + HAS_CUDA_GRAPH_SCOPE = True +except ImportError: + CudaGraphScope = None + HAS_CUDA_GRAPH_SCOPE = False + +from megatron.core.transformer.identity_op import IdentityOp +from megatron.core.transformer.module import MegatronModule +from megatron.core.transformer.spec_utils import ModuleSpec, build_module +from megatron.core.transformer.transformer_layer import TransformerLayer +from megatron.core.transformer.utils import sharded_state_dict_default +from megatron.core.utils import ( + WrappedTensor, + deprecate_inference_params, + make_viewless_tensor, +) + + +@dataclass +class HybridStackSubmodules: + """ + A class for the module specs for the MambaStack. + """ + + mamba_layer: Union[ModuleSpec, type] = IdentityOp + attention_layer: Union[ModuleSpec, type] = IdentityOp + mlp_layer: Union[ModuleSpec, type] = IdentityOp + moe_layer: Union[ModuleSpec, type] = IdentityOp + + +class HybridStack(MegatronModule): + """ + Constructor for the HybridStack class. + + Args: + config (TransformerConfig): the model configuration + submodules (MambaStackSubmodules): the submodules for the stack + residual_in_fp32 (bool, optional): whether to do residual connections + in fp32. Defaults to False. + pre_process (bool, optional): whether to include an embedding layer. + Defaults to True. + hybrid_attention_ratio (float, optional): the target ratio of attention layers to + total layers. Defaults to 0.0. + hybrid_mlp_ratio (float, optional): the target ratio of mlp layers to total + layers. Defaults to 0.0. + hybrid_override_pattern (str, optional): the hybrid layer pattern to override + with. Defaults to None. + post_layer_norm (bool, optional): whether to include a final layer norm. + Defaults to True. + post_process (bool, optional): whether to include an output layer. + Defaults to True. + device (optional): the device to use. Defaults to None. + dtype (optional): the data type to use. Defaults to None. + pg_collection (ProcessGroupCollection): the required model communication + process groups to use. + """ + + def __init__( + self, + config: TransformerConfig, + submodules: HybridStackSubmodules, + residual_in_fp32=False, + pre_process: bool = True, + hybrid_attention_ratio: float = 0.0, + hybrid_mlp_ratio: float = 0.0, + hybrid_override_pattern: str = None, + post_layer_norm: bool = True, + post_process: bool = True, + device=None, + dtype=None, + pg_collection: ProcessGroupCollection = None, + ) -> None: + super().__init__(config=config) + self.residual_in_fp32 = residual_in_fp32 + self.pre_process = pre_process + self.post_layer_norm = post_layer_norm + self.post_process = post_process + + assert pg_collection is not None, "pg_collection must be provided for MambaStack" + + self.pp_group = pg_collection.pp + self.tp_group = pg_collection.tp + + # Required for pipeline parallel schedules + self.input_tensor = None + + self.hybrid_attention_ratio = hybrid_attention_ratio + self.hybrid_mlp_ratio = hybrid_mlp_ratio + self.hybrid_override_pattern = hybrid_override_pattern + + # Customized layer allocation + # hybrid_mlp_ratio is not used in this hybrid stack. + # It is by default to be always followed by mamba or mla (i.e., mamba + MLP or MLA + MLP) + # By setting hybrid_attention_ratio, attention layers are by default to be distributed uniformly. + self.layer_type_list = self.allocate_layers( + self.config.num_layers, + self.hybrid_attention_ratio, + ) + + pp_layer_offset = 0 + if self.pp_group.size() > 1: + pp_layer_offset, self.layer_type_list = self._select_layers_for_pipeline_parallel( + self.layer_type_list + ) + + print(f"layer_type_list: {self.layer_type_list}") + + self.layers = nn.ModuleList() + for i, layer_type in enumerate(self.layer_type_list): + fp8_init_context = get_fp8_context(self.config, i + pp_layer_offset, is_init=True) + with fp8_init_context: + if layer_type == LayerSymbols.MAMBA: + layer = build_module( + submodules.mamba_layer, + config=self.config, + residual_in_fp32=residual_in_fp32, + layer_number=i + 1, + pg_collection=pg_collection, + ) + elif layer_type == LayerSymbols.ATTENTION: + # Transformer layers apply their own pp_layer_offset + layer = build_module( + submodules.attention_layer, + config=self.config, + layer_number=i + 1, + pg_collection=pg_collection, + ) + elif layer_type == LayerSymbols.MLP: + # Transformer layers apply their own pp_layer_offset + layer = build_module( + submodules.mlp_layer, + config=self.config, + layer_number=i + 1, + pg_collection=pg_collection, + ) + elif layer_type == LayerSymbols.MOE: + # Transformer layers apply their own pp_layer_offset + layer = build_module(submodules.moe_layer, config=self.config, layer_number=i + 1) + else: + assert False, "unexpected layer_type" + self.layers.append(layer) + + # Required for activation recomputation + self.num_layers_per_pipeline_rank = len(self.layers) + + if self.post_process and self.post_layer_norm: + # Final layer norm before output. + self.final_norm = TENorm( + config=self.config, + hidden_size=self.config.hidden_size, + eps=self.config.layernorm_epsilon, + ) + + def allocate_layers(self, num_layers, hybrid_attention_ratio): + layer_type_list = [] + num_attention_layers = int(num_layers // 2 * hybrid_attention_ratio) + num_mamba_layers = num_layers // 2 - num_attention_layers + num_mamba_per_attention_layer = num_mamba_layers // num_attention_layers + + if hybrid_attention_ratio <= 0.5: + base_block = [LayerSymbols.ATTENTION, LayerSymbols.MLP] + [ + LayerSymbols.MAMBA, + LayerSymbols.MLP, + ] * num_mamba_per_attention_layer + layer_type_list += base_block * num_attention_layers + layer_type_list += [LayerSymbols.MAMBA, LayerSymbols.MLP] * ( + num_mamba_layers % num_attention_layers + ) + else: + base_block = [LayerSymbols.ATTENTION, LayerSymbols.MLP] + [LayerSymbols.MAMBA, LayerSymbols.MLP] + layer_type_list += [LayerSymbols.ATTENTION, LayerSymbols.MLP] * ( + num_attention_layers - num_mamba_layers + ) + layer_type_list += base_block * num_mamba_layers + return layer_type_list + + def _select_layers_for_pipeline_parallel(self, layer_type_list): + num_layers_per_pipeline_rank = self.config.num_layers // self.pp_group.size() + + assert self.config.virtual_pipeline_model_parallel_size is None, ( + "The Mamba hybrid model does not currently support " "virtual/interleaved pipeline parallelism" + ) + + offset = self.pp_group.rank() * num_layers_per_pipeline_rank + selected_list = layer_type_list[offset : offset + num_layers_per_pipeline_rank] + + return offset, selected_list + + def set_input_tensor(self, input_tensor: Tensor): + """Set input tensor to be used instead of forward()'s input. + + When doing pipeline parallelism the input from the previous + stage comes from communication, not from the input, so the + model's forward_step_func won't have it. This function is thus + used by internal code to bypass the input provided by the + forward_step_func""" + self.input_tensor = input_tensor + + def mamba_state_shapes_per_request(self) -> Optional[Tuple[Tuple[int], Tuple[int]]]: + """ + Returns the Mamba conv and ssm states shapes per input sequence + if this block contains Mamba layers (this may not be the case with PP > 1). + """ + for layer_type, layer in zip(self.layer_type_list, self.layers): + if layer_type == LayerSymbols.MAMBA: + return layer.mamba_state_shapes_per_request() + return None + + def forward( + self, + hidden_states: Union[Tensor, WrappedTensor], + attention_mask: Tensor, + inference_context: Optional[BaseInferenceContext] = None, + rotary_pos_emb: Optional[Tensor] = None, + *, + inference_params: Optional[BaseInferenceContext] = None, + ): + """ + Forward function of the MambaStack class. + + It either returns the Loss values if labels are given or the + final hidden units + + Args: + hidden_states (Union[Tensor, WrappedTensor]): the input tensor. + Can be passed as a WrappedTensor during inference to avoid an obsolete + reference in the calling function. + attention_mask (Tensor): the attention mask. + inference_context (BaseInferenceContext): the inference parameters. + rotary_pos_emb (Tensor, optional): the rotary positional embeddings. + Defaults to None. + Returns: + Tensor: the output tensor. + """ + + inference_context = deprecate_inference_params(inference_context, inference_params) + + if not self.pre_process: + # See set_input_tensor() + hidden_states = self.input_tensor + + # Delete the obsolete reference to the initial input tensor if necessary + if isinstance(hidden_states, WrappedTensor): + hidden_states = hidden_states.unwrap() + + if inference_context and inference_context.is_static_batching(): + # NOTE(bnorick): match BaseInferenceContext attributes for + # mamba_ssm.utils.generation.BaseInferenceContext, + # this hack supports eval + inference_context.max_seqlen = inference_context.max_sequence_length + inference_context.seqlen_offset = inference_context.sequence_len_offset + + if ( + ( + ( + HAS_CUDA_GRAPH_SCOPE + and self.config.cuda_graph_impl == "local" + and CudaGraphScope.full_iteration not in self.config.cuda_graph_scope + ) + or self.config.flash_decode + ) + and inference_context + and inference_context.is_static_batching() + and not self.training + ): + current_batch_size = hidden_states.shape[1] + sequence_len_offset = torch.tensor( + [inference_context.sequence_len_offset] * current_batch_size, + dtype=torch.int32, + device="cuda", + ) + else: + sequence_len_offset = None + + # If fp8_recipe is delayed, wrap the entire pass with get_fp8_context(), + # otherwise do nothing extra at the outer level + # if we are using other fp8 recipes, then the context manager enter&exit are free + # we can wrap fp8_context within the for loop over layers, so that we can fine-grained + # control which layer will be fp8 or bf16 + use_outer_fp8_context = self.config.fp8 and self.config.fp8_recipe == Fp8Recipe.delayed + use_inner_fp8_context = self.config.fp8 and self.config.fp8_recipe != Fp8Recipe.delayed + outer_fp8_context = get_fp8_context(self.config) if use_outer_fp8_context else nullcontext() + + with outer_fp8_context: + for layer in self.layers: + inner_fp8_context = ( + get_fp8_context(self.config, layer.layer_number - 1) + if use_inner_fp8_context + else nullcontext() + ) + with inner_fp8_context: + if isinstance(layer, TransformerLayer): + hidden_states, _ = layer( + hidden_states=hidden_states, + attention_mask=attention_mask, + inference_context=inference_context, + rotary_pos_emb=rotary_pos_emb, + sequence_len_offset=sequence_len_offset, + ) + else: # MambaLayer + hidden_states = layer( + hidden_states=hidden_states, + attention_mask=attention_mask, + inference_context=inference_context, + ) + + # The attention layer (currently a simplified transformer layer) + # outputs a tuple of (hidden_states, context). Context is intended + # for cross-attention, and is not needed in our model. + if isinstance(hidden_states, tuple): + hidden_states = hidden_states[0] + + # Final layer norm. + if self.post_process and self.post_layer_norm: + hidden_states = self.final_norm(hidden_states) + + # Ensure that the tensor passed between pipeline parallel stages is + # viewless. See related notes in TransformerBlock and TransformerLayer + return make_viewless_tensor( + inp=hidden_states, requires_grad=hidden_states.requires_grad, keep_graph=True + ) + + def sharded_state_dict( + self, + prefix: str = "", + sharded_offsets: Optional[tuple] = None, + metadata: Optional[dict] = None, + ) -> ShardedStateDict: + """ + Returns a sharded state dictionary for the current object. + + This function constructs a sharded state dictionary by iterating over the layers + in the current object, computing the sharded state dictionary for each layer, + and combining the results into a single dictionary. + + Parameters: + prefix (str): The prefix to use for the state dictionary keys. + sharded_offsets (tuple): The sharded offsets to use for the state dictionary. + metadata (dict): Additional metadata to use when computing the sharded state dictionary. + + Returns: + dict: The sharded state dictionary for the current object. + """ + + sharded_state_dict = {} + layer_prefix = f"{prefix}layers." + + for local_layer_idx, layer in enumerate(self.layers): + + global_layer_offset = layer.layer_number - 1 # self.layer_number starts at 1 + state_dict_prefix = f"{layer_prefix}{local_layer_idx}." # module list index in MambaBlock + + sharded_prefix = f"{layer_prefix}{global_layer_offset}." + sharded_pp_offset = [] + + layer_sharded_state_dict = layer.sharded_state_dict( + state_dict_prefix, sharded_pp_offset, metadata + ) + + replace_prefix_for_sharding(layer_sharded_state_dict, state_dict_prefix, sharded_prefix) + + sharded_state_dict.update(layer_sharded_state_dict) + + # Add modules other than self.layers + for name, module in self.named_children(): + if not module is self.layers: + sharded_state_dict.update( + sharded_state_dict_default( + module, + f"{prefix}{name}.", + sharded_offsets, + metadata, + tp_group=self.tp_group, + ) + ) + + return sharded_state_dict diff --git a/primus/backends/megatron/core/models/hybrid/hybrid_mamba_mla_layer_specs.py b/primus/backends/megatron/core/models/hybrid/hybrid_mamba_mla_layer_specs.py new file mode 100644 index 000000000..cb801809d --- /dev/null +++ b/primus/backends/megatron/core/models/hybrid/hybrid_mamba_mla_layer_specs.py @@ -0,0 +1,119 @@ +# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +from megatron.core.extensions.transformer_engine import ( + TEColumnParallelLinear, + TEDotProductAttention, + TELayerNormColumnParallelLinear, + TELinear, + TENorm, + TERowParallelLinear, +) +from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add +from megatron.core.models.gpt.moe_module_specs import get_moe_module_spec +from megatron.core.ssm.mamba_layer import MambaLayer, MambaLayerSubmodules +from megatron.core.ssm.mamba_mixer import MambaMixer, MambaMixerSubmodules +from megatron.core.ssm.mlp_layer import MLPLayer +from megatron.core.transformer.identity_op import IdentityOp +from megatron.core.transformer.multi_latent_attention import ( + MLASelfAttention, + MLASelfAttentionSubmodules, +) + +# Import HybridStack from relative path +from primus.backends.megatron.core.models.hybrid.hybrid_block import ( + HybridStack, + HybridStackSubmodules, +) + +# Inference layers may not be available in older Megatron versions +# They're only used in hybrid_inference_stack_spec, not the training spec +try: + from megatron.core.tensor_parallel import ( + InferenceLayerNormColumnParallelLinear, + InferenceRowParallelLinear, + ) + + HAS_INFERENCE_LAYERS = True +except ImportError: + # Fallback to regular layers for inference spec + InferenceLayerNormColumnParallelLinear = TELayerNormColumnParallelLinear + InferenceRowParallelLinear = TERowParallelLinear + HAS_INFERENCE_LAYERS = False + +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.mlp import MLP, MLPSubmodules +from megatron.core.transformer.spec_utils import ModuleSpec +from megatron.core.transformer.transformer_layer import ( + TransformerLayer, + TransformerLayerSubmodules, +) + +moe = get_moe_module_spec( + use_te=True, + num_experts=8, # Can be any positive integer (must not be None). + moe_grouped_gemm=True, + moe_use_legacy_grouped_gemm=False, +) + +hybrid_stack_spec = ModuleSpec( + module=HybridStack, + submodules=HybridStackSubmodules( + mamba_layer=ModuleSpec( + module=MambaLayer, + submodules=MambaLayerSubmodules( + mixer=ModuleSpec( + module=MambaMixer, + params={ + "expand": 1, + "d_conv": 4, + }, + submodules=MambaMixerSubmodules( + in_proj=TELayerNormColumnParallelLinear, out_proj=TERowParallelLinear + ), + ), + mamba_bda=get_bias_dropout_add, + ), + ), + attention_layer=ModuleSpec( + module=TransformerLayer, + submodules=TransformerLayerSubmodules( + input_layernorm=TENorm, + self_attention=ModuleSpec( + module=MLASelfAttention, + params={"attn_mask_type": AttnMaskType.causal}, + submodules=MLASelfAttentionSubmodules( + linear_q_proj=TEColumnParallelLinear, + linear_q_down_proj=TELinear, + linear_q_up_proj=TELayerNormColumnParallelLinear, + linear_kv_down_proj=TELinear, + linear_kv_up_proj=TELayerNormColumnParallelLinear, + core_attention=TEDotProductAttention, + linear_proj=TERowParallelLinear, + q_layernorm=IdentityOp, + kv_layernorm=IdentityOp, + ), + ), + self_attn_bda=get_bias_dropout_add, + ), + ), + mlp_layer=ModuleSpec( + module=MLPLayer, + submodules=TransformerLayerSubmodules( + mlp=ModuleSpec( + module=MLP, + submodules=MLPSubmodules( + linear_fc1=TELayerNormColumnParallelLinear, linear_fc2=TERowParallelLinear + ), + ), + mlp_bda=get_bias_dropout_add, + ), + ), + moe_layer=ModuleSpec( + # TODO (rwaleffe): change this to be an "MoELayer" to work with CudaGraphs? + module=TransformerLayer, + submodules=TransformerLayerSubmodules( + pre_mlp_layernorm=TENorm, mlp=moe, mlp_bda=get_bias_dropout_add + ), + ), + ), +) diff --git a/primus/backends/megatron/megatron_pretrain_trainer.py b/primus/backends/megatron/megatron_pretrain_trainer.py index 02f97e626..97ce5e869 100644 --- a/primus/backends/megatron/megatron_pretrain_trainer.py +++ b/primus/backends/megatron/megatron_pretrain_trainer.py @@ -19,13 +19,30 @@ def train(self): from megatron.core.enums import ModelType from megatron.training import pretrain # type: ignore - from pretrain_gpt import ( # type: ignore - forward_step, - train_valid_test_datasets_provider, - ) from primus.core.utils.import_utils import get_model_provider + # Determine model type (gpt or mamba) from backend_args + model_type = getattr(self.backend_args, "model_type", "gpt") + log_rank_0(f"-detected model_type: {model_type}") + + # Import the appropriate training components based on model_type + if model_type == "mamba": + from pretrain_mamba import ( # type: ignore + forward_step, + train_valid_test_datasets_provider, + ) + + log_rank_0("Using Mamba model provider and training components") + else: + from pretrain_gpt import ( # type: ignore + forward_step, + train_valid_test_datasets_provider, + ) + + log_rank_0("Using GPT model provider and training components") + + # Configure training components if hasattr(train_valid_test_datasets_provider, "is_distributed"): train_valid_test_datasets_provider.is_distributed = True @@ -49,9 +66,17 @@ def train(self): if "store" in sig.parameters: kwargs["store"] = store + # Get model provider with correct model_type + # Only pass model_type if it's not the default to maintain compatibility + if model_type != "gpt": + model_provider = get_model_provider(model_type=model_type) + else: + model_provider = get_model_provider() + log_rank_0(f"-model_provider: {model_provider}") + wrapped_pretrain( train_valid_test_datasets_provider, - get_model_provider(), + model_provider, ModelType.encoder_or_decoder, forward_step, **kwargs, diff --git a/primus/configs/models/megatron/deepseek_v2_base.yaml b/primus/configs/models/megatron/deepseek_v2_base.yaml index 6d402635d..f4d645c17 100755 --- a/primus/configs/models/megatron/deepseek_v2_base.yaml +++ b/primus/configs/models/megatron/deepseek_v2_base.yaml @@ -34,6 +34,7 @@ moe_aux_loss_coeff: 0.001 # aux_loss_alpha # rotary rotary_base: 10000 rotary_scaling_factor: 40.0 # float +original_max_position_embeddings: 4096 # int, original context length before YaRN scaling mscale: 0.707 # float mscale_all_dim: 0.707 # float diff --git a/primus/configs/models/megatron/deepseek_v3_base.yaml b/primus/configs/models/megatron/deepseek_v3_base.yaml index 421b99361..c0a5c77af 100755 --- a/primus/configs/models/megatron/deepseek_v3_base.yaml +++ b/primus/configs/models/megatron/deepseek_v3_base.yaml @@ -37,6 +37,7 @@ moe_aux_loss_coeff: 0.001 # aux_loss_alpha # rotary rotary_base: 10000 rotary_scaling_factor: 40.0 # float +original_max_position_embeddings: 4096 # int, original context length before YaRN scaling mscale: 1.0 # float mscale_all_dim: 1.0 # float diff --git a/primus/configs/models/megatron/hybrid_model_base.yaml b/primus/configs/models/megatron/hybrid_model_base.yaml new file mode 100644 index 000000000..f6e9ee113 --- /dev/null +++ b/primus/configs/models/megatron/hybrid_model_base.yaml @@ -0,0 +1,21 @@ +extends: + - language_model.yaml + +# Mamba layer configuration +mamba_state_dim: 128 +mamba_head_dim: 64 +mamba_num_groups: 8 +mamba_num_heads: null +mamba_expand: 2 +mamba_d_conv: 4 +disable_mamba_mem_eff_path: false + +# Hybrid model configuration +is_hybrid_model: false # bool +hybrid_attention_ratio: 0.0 # float range [0,0, 1.0] +hybrid_mlp_ratio: 0.0 # float range [0,0, 1.0] +hybrid_override_pattern: null # str + +# MTP +mtp_num_layers: null # int +mtp_loss_scaling_factor: 0.1 # float diff --git a/primus/configs/models/megatron/language_model.yaml b/primus/configs/models/megatron/language_model.yaml index 8c6f24e48..9d2929c2a 100755 --- a/primus/configs/models/megatron/language_model.yaml +++ b/primus/configs/models/megatron/language_model.yaml @@ -8,6 +8,7 @@ extends: # model architecture use_legacy_models: false deprecated_use_mcore_models: false +model_type: gpt # gpt or mamba num_layers: 24 encoder_num_layers: null decoder_num_layers: null @@ -21,6 +22,7 @@ num_query_groups: null add_position_embedding: false position_embedding_type: learned_absolute max_position_embeddings: null +original_max_position_embeddings: null untie_embeddings_and_output_weights: true ffn_hidden_size: null @@ -99,10 +101,6 @@ rotary_scaling_factor: 1.0 # float mscale: 1.0 # float mscale_all_dim: 1.0 # float -# MTP -mtp_num_layers: null # int -mtp_loss_scaling_factor: 0.1 # float - # MoE related num_experts: null moe_layer_freq: 1 # int diff --git a/primus/configs/models/megatron/mamba_370M.yaml b/primus/configs/models/megatron/mamba_370M.yaml new file mode 100644 index 000000000..6665da3af --- /dev/null +++ b/primus/configs/models/megatron/mamba_370M.yaml @@ -0,0 +1,16 @@ +bases: + - mamba_base.yaml + +# Mamba 370M configuration +model_type: mamba # CRITICAL: Mamba models must use mamba model type +tokenizer_type: GPT2BPETokenizer +vocab_size: 50257 + +# Model size parameters +num_layers: 48 +hidden_size: 1024 +num_attention_heads: 16 # Required by Megatron validation, even for pure Mamba models +ffn_hidden_size: null +mamba_state_dim: 16 +mamba_head_dim: 64 +mamba_num_groups: 8 diff --git a/primus/configs/models/megatron/mamba_base.yaml b/primus/configs/models/megatron/mamba_base.yaml new file mode 100644 index 000000000..d52fe6db2 --- /dev/null +++ b/primus/configs/models/megatron/mamba_base.yaml @@ -0,0 +1,36 @@ +bases: + - language_model.yaml + +# Mamba-specific configuration +# Note: Mamba-specific parameters (spec, is_hybrid_model, mamba_state_dim, etc.) +# must be set in the pretrain config overrides, not here + +model_type: mamba +use_legacy_models: false + +# Position embeddings - Mamba typically doesn't use position embeddings +position_embedding_type: rope +use_rotary_position_embeddings: false + +# Tokenizer (should be set in specific model configs) +tokenizer_type: HuggingFaceTokenizer +tokenizer_model: null + +# Standard transformer settings that may be used by hybrid models +is_hybrid_model: false +attention_dropout: 0.0 +hidden_dropout: 0.0 + +# Embeddings +untie_embeddings_and_output_weights: false + +# Other settings +apply_residual_connection_post_layernorm: false +add_bias_linear: false +swiglu: false + +# Normalization +norm_epsilon: 1.0e-5 + +# Initialization +init_method_std: 0.02 diff --git a/primus/configs/models/megatron/zebra_llama_1B.yaml b/primus/configs/models/megatron/zebra_llama_1B.yaml new file mode 100644 index 000000000..d1afa9531 --- /dev/null +++ b/primus/configs/models/megatron/zebra_llama_1B.yaml @@ -0,0 +1,42 @@ +bases: + - mamba_base.yaml + +# Zebra Llama 8B configuration +model_type: mamba # CRITICAL: Hybrid models must use mamba model type +tokenizer_type: HuggingFaceTokenizer +tokenizer_model: meta-llama/Llama-3.2-1B + +# Model size parameters +num_layers: 32 +hidden_size: 2048 +ffn_hidden_size: 8192 + +# Mamba parameters +is_hybrid_model: true +hybrid_attention_ratio: 0.25 +mamba_state_dim: 64 +mamba_head_dim: 64 +mamba_num_groups: 8 + +# MLA parameters +# Disable standard GQA - MLA uses its own compression via LoRA +group_query_attention: false +swiglu: true +num_query_groups: null +multi_latent_attention: true +num_attention_heads: 32 +q_lora_rank: 1344 # Query LoRA rank +kv_lora_rank: 128 # Key-Value LoRA rank +qk_head_dim: 32 # Query-Key head dimension +qk_pos_emb_head_dim: 32 # Positional embedding head dimension +v_head_dim: 64 # Value head dimension +rotary_scaling_factor: 1.0 +mscale: 1.0 +mscale_all_dim: 1.0 + +# MLA uses its own internal positional encoding +rotary_base: 500000 +position_embedding_type: none +add_position_embedding: true +use_rotary_position_embeddings: false +max_position_embeddings: 131072 diff --git a/primus/configs/models/megatron/zebra_llama_3B.yaml b/primus/configs/models/megatron/zebra_llama_3B.yaml new file mode 100644 index 000000000..23090841f --- /dev/null +++ b/primus/configs/models/megatron/zebra_llama_3B.yaml @@ -0,0 +1,43 @@ +bases: + - mamba_base.yaml + +# Zebra Llama 8B configuration +model_type: mamba # CRITICAL: Hybrid models must use mamba model type +tokenizer_type: HuggingFaceTokenizer +tokenizer_model: meta-llama/Llama-3.2-1B + +# Model size parameters +num_layers: 56 +hidden_size: 3072 +ffn_hidden_size: 8192 +normalization: "RMSNorm" + +# Mamba parameters +is_hybrid_model: true +hybrid_attention_ratio: 0.25 +mamba_state_dim: 128 +mamba_head_dim: 128 +mamba_num_groups: 8 + +# MLA parameters +# Disable standard GQA - MLA uses its own compression via LoRA +group_query_attention: false +swiglu: true +num_query_groups: null +multi_latent_attention: true +num_attention_heads: 24 +q_lora_rank: 1536 # Query LoRA rank +kv_lora_rank: 128 # Key-Value LoRA rank +qk_head_dim: 64 # Query-Key head dimension +qk_pos_emb_head_dim: 64 # Positional embedding head dimension +v_head_dim: 128 # Value head dimension +rotary_scaling_factor: 1.0 +mscale: 1.0 +mscale_all_dim: 1.0 + +# MLA uses its own internal positional encoding +rotary_base: 500000 +position_embedding_type: none +add_position_embedding: true +use_rotary_position_embeddings: false +original_max_position_embeddings: 131072 diff --git a/primus/configs/models/megatron/zebra_llama_8B.yaml b/primus/configs/models/megatron/zebra_llama_8B.yaml new file mode 100644 index 000000000..0237a652d --- /dev/null +++ b/primus/configs/models/megatron/zebra_llama_8B.yaml @@ -0,0 +1,43 @@ +bases: + - mamba_base.yaml + +# Zebra Llama 8B configuration +model_type: mamba # CRITICAL: Hybrid models must use mamba model type +tokenizer_type: HuggingFaceTokenizer +tokenizer_model: meta-llama/Llama-3.2-1B + +# Model size parameters +num_layers: 64 +hidden_size: 4096 +ffn_hidden_size: 14436 +normalization: "RMSNorm" + +# Mamba parameters +is_hybrid_model: true +hybrid_attention_ratio: 0.25 +mamba_state_dim: 128 +mamba_head_dim: 128 +mamba_num_groups: 8 + +# MLA parameters +# Disable standard GQA - MLA uses its own compression via LoRA +group_query_attention: false +swiglu: true +num_query_groups: null +multi_latent_attention: true +num_attention_heads: 32 +q_lora_rank: 2048 # Query LoRA rank +kv_lora_rank: 160 # Key-Value LoRA rank +qk_head_dim: 64 # Query-Key head dimension +qk_pos_emb_head_dim: 64 # Positional embedding head dimension +v_head_dim: 128 # Value head dimension +rotary_scaling_factor: 1.0 +mscale: 1.0 +mscale_all_dim: 1.0 + +# MLA uses its own internal positional encoding +rotary_base: 500000 +position_embedding_type: none +add_position_embedding: true +use_rotary_position_embeddings: false +original_max_position_embeddings: 131072 diff --git a/primus/configs/modules/megatron/trainer_base.yaml b/primus/configs/modules/megatron/trainer_base.yaml index 76be0c7ba..ba89bd7d1 100755 --- a/primus/configs/modules/megatron/trainer_base.yaml +++ b/primus/configs/modules/megatron/trainer_base.yaml @@ -305,17 +305,17 @@ rerun_mode: disabled # str: 'disabled', 'validate_results', 'report_stats' # Experimental features enable_experimental: false -# Hybrid model configuration -hybrid_attention_ratio: 0.0 # float range [0,0, 1.0] -hybrid_mlp_ratio: 0.0 # float range [0,0, 1.0] -hybrid_override_pattern: null # str - -# Mamba layer configuration -mamba_state_dim: 128 -mamba_head_dim: 64 -mamba_num_groups: 8 -mamba_num_heads: null -disable_mamba_mem_eff_path: false +# # Hybrid model configuration +# hybrid_attention_ratio: 0.0 # float range [0,0, 1.0] +# hybrid_mlp_ratio: 0.0 # float range [0,0, 1.0] +# hybrid_override_pattern: null # str + +# # Mamba layer configuration +# mamba_state_dim: 128 +# mamba_head_dim: 64 +# mamba_num_groups: 8 +# mamba_num_heads: null +# disable_mamba_mem_eff_path: false # Args of precision-aware optimizer use_precision_aware_optimizer: false @@ -405,7 +405,6 @@ indexer_log_interval: 1000 enable_ft_package: false calc_ft_timeouts: false run_workload_inspector_server: false -is_hybrid_model: false heterogeneous_layers_config_path: null heterogeneous_layers_config_encoded_json: null diff --git a/primus/core/utils/import_utils.py b/primus/core/utils/import_utils.py index 2ccd8ebed..34c2de16d 100644 --- a/primus/core/utils/import_utils.py +++ b/primus/core/utils/import_utils.py @@ -34,25 +34,42 @@ def lazy_import(paths, symbol, log_prefix="[Primus]"): raise ImportError(f"{log_prefix} {symbol} not found in any of: {paths}") -def get_model_provider(): +def get_model_provider(model_type="gpt"): """ - Resolve model_provider across Megatron versions. + Resolve model_provider across Megatron versions and model types. - - New: model_provider + gpt_builder + Args: + model_type (str): Type of model - 'gpt' or 'mamba'. Defaults to 'gpt'. + + - New: model_provider + gpt_builder/mamba_builder - Mid: model_provider only - - Old: pretrain_gpt.model_provider + - Old: pretrain_gpt.model_provider / pretrain_mamba.model_provider """ # Try to import model_provider - model_provider = lazy_import( - ["model_provider", "pretrain_gpt"], "model_provider", log_prefix="[Primus][MegatronCompat]" - ) + if model_type == "mamba": + model_provider = lazy_import( + ["model_provider", "pretrain_mamba"], "model_provider", log_prefix="[Primus][MegatronCompat]" + ) + # Try to import mamba_builder (for Mamba models) + try: + mamba_builder = lazy_import( + ["mamba_builders"], "mamba_builder", log_prefix="[Primus][MegatronCompat]" + ) + return partial(model_provider, mamba_builder) + except ImportError: + return model_provider + else: + # Default GPT behavior + model_provider = lazy_import( + ["model_provider", "pretrain_gpt"], "model_provider", log_prefix="[Primus][MegatronCompat]" + ) - # Try to import gpt_builder (only exists in newer versions) - try: - gpt_builder = lazy_import(["gpt_builders"], "gpt_builder", log_prefix="[Primus][MegatronCompat]") - return partial(model_provider, gpt_builder) - except ImportError: - return model_provider + # Try to import gpt_builder (only exists in newer versions) + try: + gpt_builder = lazy_import(["gpt_builders"], "gpt_builder", log_prefix="[Primus][MegatronCompat]") + return partial(model_provider, gpt_builder) + except ImportError: + return model_provider def get_custom_fsdp(): diff --git a/primus/modules/trainer/megatron/pre_trainer.py b/primus/modules/trainer/megatron/pre_trainer.py index 1c8539dd5..d34788b04 100644 --- a/primus/modules/trainer/megatron/pre_trainer.py +++ b/primus/modules/trainer/megatron/pre_trainer.py @@ -242,6 +242,20 @@ def forward_step(self, data_iterator, model: GPTModel, return_schedule_plan=Fals assert ( args.overlap_moe_expert_parallel_comm ), "overlap_moe_expert_parallel_comm must be enabled to return the schedule plan" + + # Schedule plan building is only supported for GPT models + # Check if this is a Mamba model + unwrapped_model = model + while hasattr(unwrapped_model, "module"): + unwrapped_model = unwrapped_model.module + model_class_name = unwrapped_model.__class__.__name__ + + if "Mamba" in model_class_name: + raise NotImplementedError( + "Schedule plan building is not supported for Mamba models. " + "Please disable overlap_moe_expert_parallel_comm for Mamba." + ) + if args.patch_moe_overlap: assert ( not args.delay_wgrad_compute @@ -267,8 +281,21 @@ def forward_step(self, data_iterator, model: GPTModel, return_schedule_plan=Fals ) return schedule_plan, partial(self.loss_func, loss_mask) else: - output_tensor = model( - tokens, position_ids, attention_mask, labels=labels, loss_mask=loss_mask - ) + # Check if model supports loss_mask parameter + # MambaModel doesn't accept loss_mask, but GPTModel does + # Unwrap the model to get the actual model class + unwrapped_model = model + while hasattr(unwrapped_model, "module"): + unwrapped_model = unwrapped_model.module + model_class_name = unwrapped_model.__class__.__name__ + + if "Mamba" in model_class_name: + # MambaModel doesn't accept loss_mask parameter + output_tensor = model(tokens, position_ids, attention_mask, labels=labels) + else: + # GPTModel and other models accept loss_mask parameter + output_tensor = model( + tokens, position_ids, attention_mask, labels=labels, loss_mask=loss_mask + ) return output_tensor, partial(self.loss_func, loss_mask) diff --git a/primus/modules/trainer/megatron/trainer.py b/primus/modules/trainer/megatron/trainer.py index 5db59188b..8f0c9f70e 100644 --- a/primus/modules/trainer/megatron/trainer.py +++ b/primus/modules/trainer/megatron/trainer.py @@ -483,8 +483,10 @@ def update_primus_config( if args.iterations_to_skip is None: args.iterations_to_skip = [] - # support moe_freq_type - if isinstance(args.moe_layer_freq, str): + # support moe_freq_type - ensure moe_layer_freq has a default value + if not hasattr(args, "moe_layer_freq"): + args.moe_layer_freq = 1 + elif isinstance(args.moe_layer_freq, str): try: args.moe_layer_freq = eval(args.moe_layer_freq) except Exception: @@ -496,11 +498,35 @@ def update_primus_config( args.valid_data_path = None args.test_data_path = None + # Determine model type (gpt or mamba) + model_type = getattr(args, "model_type", "gpt") + log_rank_0(f"-detected model_type: {model_type}") + + # Ensure required attributes have safe defaults if missing from config + if not hasattr(args, "final_logit_softcapping"): + args.final_logit_softcapping = None + if not hasattr(args, "router_logit_softcapping"): + args.router_logit_softcapping = None + + # Only pass model_type parameter when it's "mamba" to maintain backward compatibility + # with main branch behavior for "gpt" (default) case if args.final_logit_softcapping is not None and args.final_logit_softcapping > 0.0: log_rank_0(f"-enable final_logit_softcapping: {args.final_logit_softcapping}") - self.model_provider = functools.partial(primus_model_provider, get_model_provider()) + if model_type == "mamba": + self.model_provider = functools.partial( + primus_model_provider, get_model_provider(model_type=model_type) + ) + else: + self.model_provider = functools.partial(primus_model_provider, get_model_provider()) else: - self.model_provider = get_model_provider() + if model_type == "mamba": + log_rank_0(f"-getting model provider for model_type={model_type}") + model_provider = get_model_provider(model_type=model_type) + log_rank_0(f"-model_provider: {model_provider}") + self.model_provider = model_provider + else: + # For "gpt" (default), call without arguments to match main branch behavior + self.model_provider = get_model_provider() if args.router_logit_softcapping is not None and args.router_logit_softcapping > 0.0: log_rank_0(f"-enable router_logit_softcapping: {args.router_logit_softcapping}") @@ -867,6 +893,8 @@ def setup_model_and_optimizer( log_rank_0(f"use te backend...") log_rank_0(f"-run get_model") + log_rank_0(f"-model_provider_func: {model_provider_func}") + log_rank_0(f"-model_type: {model_type}") model = get_model(model_provider_func, model_type) log_rank_0(model) # get_megatron_optimizer will use the ddp_config From 365758c254fcc1b6aba1ef4a9879f6538bef63d1 Mon Sep 17 00:00:00 2001 From: clairesonglee Date: Fri, 27 Feb 2026 05:23:49 +0000 Subject: [PATCH 06/11] tune hybrid model mi300x configs --- examples/megatron/configs/MI300X/zebra_llama_1B-pretrain.yaml | 4 ++-- examples/megatron/configs/MI300X/zebra_llama_3B-pretrain.yaml | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/examples/megatron/configs/MI300X/zebra_llama_1B-pretrain.yaml b/examples/megatron/configs/MI300X/zebra_llama_1B-pretrain.yaml index d2327bab8..ca65bb754 100644 --- a/examples/megatron/configs/MI300X/zebra_llama_1B-pretrain.yaml +++ b/examples/megatron/configs/MI300X/zebra_llama_1B-pretrain.yaml @@ -21,8 +21,8 @@ modules: log_avg_reset_interval: 50 train_iters: 100 - micro_batch_size: 2 - global_batch_size: 16 + micro_batch_size: 8 + global_batch_size: 64 seq_length: 8192 max_position_embeddings: 8192 diff --git a/examples/megatron/configs/MI300X/zebra_llama_3B-pretrain.yaml b/examples/megatron/configs/MI300X/zebra_llama_3B-pretrain.yaml index 4daeb25e6..b8019be76 100644 --- a/examples/megatron/configs/MI300X/zebra_llama_3B-pretrain.yaml +++ b/examples/megatron/configs/MI300X/zebra_llama_3B-pretrain.yaml @@ -21,8 +21,8 @@ modules: log_avg_reset_interval: 50 train_iters: 100 - micro_batch_size: 2 - global_batch_size: 16 + micro_batch_size: 4 + global_batch_size: 32 seq_length: 8192 max_position_embeddings: 8192 From 06d8e1e86852002f5bcb1f784784d6bbf3317ec9 Mon Sep 17 00:00:00 2001 From: clairesonglee Date: Fri, 27 Feb 2026 05:33:34 +0000 Subject: [PATCH 07/11] tune hybrid model mi355x configs --- examples/megatron/configs/MI355X/zebra_llama_1B-pretrain.yaml | 4 ++-- examples/megatron/configs/MI355X/zebra_llama_3B-pretrain.yaml | 4 ++-- examples/megatron/configs/MI355X/zebra_llama_8B-pretrain.yaml | 4 ++-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/examples/megatron/configs/MI355X/zebra_llama_1B-pretrain.yaml b/examples/megatron/configs/MI355X/zebra_llama_1B-pretrain.yaml index 4eaf5102c..67647da43 100644 --- a/examples/megatron/configs/MI355X/zebra_llama_1B-pretrain.yaml +++ b/examples/megatron/configs/MI355X/zebra_llama_1B-pretrain.yaml @@ -21,8 +21,8 @@ modules: log_avg_reset_interval: 50 train_iters: 100 - micro_batch_size: 16 - global_batch_size: 128 + micro_batch_size: 12 + global_batch_size: 96 seq_length: 8192 max_position_embeddings: 8192 diff --git a/examples/megatron/configs/MI355X/zebra_llama_3B-pretrain.yaml b/examples/megatron/configs/MI355X/zebra_llama_3B-pretrain.yaml index 4daeb25e6..19006bd72 100644 --- a/examples/megatron/configs/MI355X/zebra_llama_3B-pretrain.yaml +++ b/examples/megatron/configs/MI355X/zebra_llama_3B-pretrain.yaml @@ -21,8 +21,8 @@ modules: log_avg_reset_interval: 50 train_iters: 100 - micro_batch_size: 2 - global_batch_size: 16 + micro_batch_size: 7 + global_batch_size: 56 seq_length: 8192 max_position_embeddings: 8192 diff --git a/examples/megatron/configs/MI355X/zebra_llama_8B-pretrain.yaml b/examples/megatron/configs/MI355X/zebra_llama_8B-pretrain.yaml index ff274291e..136ac9d3d 100644 --- a/examples/megatron/configs/MI355X/zebra_llama_8B-pretrain.yaml +++ b/examples/megatron/configs/MI355X/zebra_llama_8B-pretrain.yaml @@ -21,8 +21,8 @@ modules: log_avg_reset_interval: 50 train_iters: 100 - micro_batch_size: 16 - global_batch_size: 128 + micro_batch_size: 4 + global_batch_size: 32 seq_length: 8192 max_position_embeddings: 8192 From fadaeb1e958e32ba9b92ba8c40e56d06ec20e78c Mon Sep 17 00:00:00 2001 From: Anshu Raina Date: Mon, 2 Mar 2026 14:55:58 -0800 Subject: [PATCH 08/11] Expand projection.md with memory projection and performance details. (#578) Expand projection.md with memory projection and performance details. --- docs/projection.md | 661 +++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 609 insertions(+), 52 deletions(-) diff --git a/docs/projection.md b/docs/projection.md index 5c5212026..9af860731 100644 --- a/docs/projection.md +++ b/docs/projection.md @@ -1,39 +1,64 @@ -# Performance Projection +# Projection -Primus includes a performance projection tool that benchmarks transformer layers on a single node and projects training iteration times to multi-node configurations. +Primus includes projection tools that estimate **memory requirements** and **training performance** for large-scale distributed training without requiring the full target cluster. Two projection modes are available: -- **User-facing entry**: `primus-cli … -- projection performance [options]` -- **Implementation entrypoint**: `primus/cli/subcommands/projection.py` -- **Core logic**: `primus/core/projection/performance_projection/projection.py` - -## Overview - -The performance projection tool: - -1. **Benchmarks** transformer layers on a single node to measure forward/backward pass times -2. **Simulates** pipeline parallelism scheduling (including zero-bubble optimization) -3. **Projects** performance to multi-node configurations by modeling: - - Data Parallelism (DP) scaling - - Gradient AllReduce communication overhead - - Expert Parallelism (EP) All-to-All communication overhead - - Inter-node vs intra-node communication differences +| Mode | Command | What it does | +|------|---------|--------------| +| **Memory** | `projection memory` | Estimates per-GPU memory (parameters, optimizer, activations) using analytical formulas | +| **Performance** | `projection performance` | Benchmarks layers on 1 node, then projects training time to multi-node clusters | -This allows you to estimate training performance on larger clusters without actually running on them. +- **User-facing entry**: `primus-cli … -- projection {memory,performance} [options]` +- **Implementation entrypoint**: `primus/cli/subcommands/projection.py` +- **Core logic**: + - Memory: `primus/core/projection/memory_projection/projection.py` + - Performance: `primus/core/projection/performance_projection/projection.py` + +--- + +## Table of Contents + +1. [Quick Start](#quick-start) +2. [Command Syntax](#command-syntax) +3. [Memory Projection](#memory-projection) + - [Overview](#memory-overview) + - [Architecture](#memory-architecture) + - [Parameter Estimation](#parameter-estimation) + - [Param + Optimizer Memory](#param--optimizer-memory) + - [Activation Memory](#activation-memory) + - [Pipeline Schedule Memory Scaling](#pipeline-schedule-memory-scaling) + - [Recomputation Support](#recomputation-support) + - [Memory Formulas Reference](#memory-formulas-reference) +4. [Performance Projection](#performance-projection) + - [Overview](#performance-overview) + - [How It Works](#how-it-works) + - [Scaling Mechanisms](#scaling-mechanisms) + - [Communication Modeling](#communication-modeling) + - [Pipeline Schedule Simulator](#pipeline-schedule-simulator) + - [Overall Performance Prediction Flow](#overall-performance-prediction-flow) +5. [Example Output](#example-output) +6. [Assumptions and Limitations](#assumptions-and-limitations) +7. [Tips](#tips) + +--- ## Quick Start -Run a basic performance projection for the minimum required nodes: +### Memory Projection + +Estimate per-GPU memory for a model configuration (no GPU needed for estimation, but the CLI currently requires torch distributed init): ```bash export NNODES=1 export HSA_NO_SCRATCH_RECLAIM=1 bash runner/primus-cli direct --script primus/cli/main.py -- \ - projection performance \ + projection memory \ --config examples/megatron/configs/MI300X/deepseek_v2_lite-BF16-pretrain.yaml ``` -Project performance to a specific number of nodes: +### Performance Projection + +Benchmark on 1 node and project to a target cluster: ```bash export NNODES=1 @@ -45,26 +70,32 @@ bash runner/primus-cli direct --script primus/cli/main.py -- \ --target-nodes 4 ``` +--- + ## Command Syntax ```bash -primus-cli [global-options] [mode-args] -- projection performance [options] +primus-cli [global-options] [mode-args] -- projection {memory,performance} [options] ``` -### Options +### Common Options | Option | Type | Description | |--------|------|-------------| | `--config` | string | Path to the Primus YAML configuration file (required) | + +### Performance-Only Options + +| Option | Type | Description | +|--------|------|-------------| | `--target-nodes` | int | Target number of nodes for projection. Defaults to minimum required by parallelism config | | `--hardware-config` | string | Path to YAML file with custom hardware parameters for communication modeling | ### Parallelism Overrides -You can override parallelism settings from the config file using environment variables or CLI overrides: +You can override parallelism settings from the config file: ```bash -# Using environment variables export PRIMUS_TP=1 export PRIMUS_PP=3 export PRIMUS_EP=8 @@ -75,30 +106,390 @@ bash runner/primus-cli direct --script primus/cli/main.py -- \ --target-nodes 6 ``` -## How It Works +--- + +## Memory Projection + + +### Overview + +The memory projection estimates **per-GPU memory** usage by analytically computing: + +1. **Parameter memory** β€” model weights stored on this GPU +2. **Optimizer state memory** β€” optimizer first/second moments, sharded across DP ranks +3. **Activation memory** β€” intermediate tensors stored for the backward pass + +It uses a hierarchical profiler system that mirrors the model's module structure, computing each component's contribution bottom-up. + + +### Architecture + +``` +LanguageModelProfiler +β”œβ”€β”€ EmbeddingProfiler β€” vocab embeddings (stage 0 only) +β”œβ”€β”€ DenseTransformerLayerProfiler β€” for non-MoE layers +β”‚ β”œβ”€β”€ LayerNormProfiler (Γ—3) β€” pre-attn, pre-MLP, post-MLP +β”‚ β”œβ”€β”€ AttentionProfiler β€” QKV projections + attention +β”‚ β”œβ”€β”€ ResidualAddProfiler (Γ—2) β€” skip connections +β”‚ └── DenseMLPProfiler β€” standard SwiGLU/FFN +β”œβ”€β”€ MoETransformerLayerProfiler β€” for MoE layers +β”‚ β”œβ”€β”€ LayerNormProfiler (Γ—3) +β”‚ β”œβ”€β”€ AttentionProfiler +β”‚ β”œβ”€β”€ ResidualAddProfiler (Γ—2) +β”‚ β”œβ”€β”€ RouterProfiler β€” expert routing logits +β”‚ └── MoEMLPProfiler β€” routed experts + shared expert +β”œβ”€β”€ LayerNormProfiler β€” final layer norm (last stage only) +β”œβ”€β”€ OutputLayerProfiler β€” language model head (last stage only) +└── LossProfiler β€” cross-entropy loss (last stage only) +``` + +Each profiler implements two key methods: +- `estimated_num_params(rank)` β€” parameter count (total if `rank=None`, per-GPU if rank given) +- `estimated_activation_memory(batch_size, seq_len)` β€” activation bytes for one microbatch + + +### Parameter Estimation + +Parameters are computed per component and summed across all layers assigned to this GPU's pipeline stage. + +#### Layer Assignment + +Layers are distributed across `PP Γ— VPP` virtual stages. Each physical PP rank hosts `VPP` virtual stages in an interleaved pattern: + +``` +PP rank 0 β†’ virtual stages 0, PP, 2Γ—PP, ... +PP rank 1 β†’ virtual stages 1, PP+1, 2Γ—PP+1, ... +``` + +#### Per-Component Parameter Formulas + +| Component | Formula | Notes | +|-----------|---------|-------| +| **Embedding** | `V Γ— H` | `V` = padded vocab size, `H` = hidden size | +| **LayerNorm** | `2 Γ— H` | gamma + beta per LayerNorm | +| **Attention (standard)** | `2 Γ— HΒ² Γ— (1 + G/A) Γ— P` | `A` = num heads, `G` = num KV groups, `P` = proj ratio | +| **Attention (MLA)** | `q_term + kv_term + pos + out` | DeepSeek-style multi-latent attention | +| **Dense MLP (SwiGLU)** | `3 Γ— H Γ— FFN` | gate, up, down projections | +| **Dense MLP (standard)** | `2 Γ— H Γ— FFN` | up, down projections | +| **Router** | `H Γ— NE` | `NE` = number of experts | +| **MoE MLP** | `NE/EP Γ— n_proj Γ— H Γ— FFN_e + shared` | Expert params sharded by EP | +| **Output Layer** | `V Γ— H` | May share weights with embedding | + +Where: +- `H` = `hidden_size` +- `V` = `padded_vocab_size` +- `FFN` = `ffn_hidden_size` (dense MLP intermediate dimension) +- `FFN_e` = `moe_ffn_hidden_size` (per-expert intermediate dimension) +- `NE` = `num_experts` +- `EP` = `expert_model_parallel_size` +- `n_proj` = 3 for SwiGLU, 2 for standard FFN + +### Param + Optimizer Memory + +The total bytes per parameter depends on the training precision and optimizer sharding: + +``` +bytes_per_param = weight_bytes + gradient_bytes + optimizer_bytes + +Where: + weight_bytes = 2 (BF16 weights) + gradient_bytes = 2 (BF16 gradients) + optimizer_bytes = 10/DP (FP32 master weights + Adam m + Adam v, sharded across DP) + = (2 + 4 + 4) / DP +``` + +**DP calculation:** + +``` +DP = world_size / (EP Γ— PP) +``` + +Note: TP and CP are not divided out because all TP/CP ranks within a DP group share the same optimizer partition. + +**Total param + optimizer memory per GPU:** + +``` +param_optimizer_memory = params_on_this_gpu Γ— bytes_per_param +``` + +### Activation Memory -### 1. Configuration Reduction +Activation memory is the memory needed to store intermediate tensors for the backward pass. Each component estimates what it stores for backward. + +#### Base Tensor (sbh) + +The fundamental building block is the hidden state tensor: + +``` +sbh = MBS Γ— (S / TP / CP) Γ— H Γ— 2 bytes (BF16) +``` + +Where `MBS` = micro batch size, `S` = sequence length. + +#### Per-Component Activation Formulas + +##### LayerNorm +Stores its input for backward: +``` +act = sbh = MBS Γ— S/(TPΓ—CP) Γ— H Γ— 2 +``` + +##### Residual Add +Stores the residual for backward: +``` +act = sbh +``` + +##### Router +Stores hidden states for routing weight gradients: +``` +act = sbh +``` + +##### Attention (standard, Flash Attention) + +Stores Q, K, V, output, and logsumexp for Flash Attention backward: + +```python +tokens_per_rank = MBS Γ— S / (TP Γ— CP) + +# activation width = Q + K + V + output + softmax stats +Q_width = kv_channels Γ— num_heads # e.g. 128 Γ— 64 = 8192 +KV_width = kv_channels Γ— num_kv_groups # e.g. 128 Γ— 1 = 128 (MQA) +output_width = hidden_size # 8192 +softmax_width = Q_width (with Flash Attention) # 8192 + +total_width = Q_width + 2Γ—KV_width + output_width + softmax_width +act = tokens_per_rank Γ— total_width Γ— 2 (BF16) +``` + +For MQA with 64 heads and 1 KV group: `Q(256MB) + K(4MB) + V(4MB) + O(256MB) + LSE(4MB) β‰ˆ 0.51 GB` + +##### Dense MLP (SwiGLU) + +For the SwiGLU computation `output = down_proj(silu(gate_proj(x)) βŠ™ up_proj(x))`, stores: + +```python +tokens = MBS Γ— S / (TP Γ— CP) + +# SwiGLU stores gate, up, and hidden (siluΓ—up) for backward +intermediate = 2 Γ— tokens Γ— FFN Γ— 2 # gate_proj + up_proj outputs (BF16) +activation = tokens Γ— FFN Γ— 2 # silu(gate) βŠ™ up (input to down_proj) +output = tokens Γ— H Γ— 2 # down_proj output + +act = intermediate + activation + output + = tokens Γ— (3Γ—FFN + H) Γ— 2 +``` + +##### MoE MLP + +For MoE, each token is routed to `topk` experts, duplicating the activation: + +```python +tokens = MBS Γ— S / (TP Γ— CP) +topk_tokens = tokens Γ— topk # total token-expert pairs + +# Routed experts: same SwiGLU structure per token-expert pair +routed_act = topk_tokens Γ— (3Γ—FFN_e + H) Γ— 2 + +# Shared expert (if configured): processes ALL tokens +shared_act = tokens Γ— (3Γ—FFN_e + H) Γ— 2 # same SwiGLU, one copy + +act = routed_act + shared_act + = tokens Γ— (topk + N_shared) Γ— (H + 3Γ—FFN_e) Γ— 2 +``` + +Where: +- `topk` = `moe_router_topk` (experts activated per token) +- `FFN_e` = `moe_ffn_hidden_size` (per-expert FFN intermediate dimension) +- `N_shared` = 1 if `moe_shared_expert_intermediate_size` is set, else 0 + +**Example (MoE 4.5T, MBS=4, S=16384, CP=4, H=8192, FFN_e=2048, topk=36):** +``` +tokens = 4 Γ— 16384/4 = 16,384 +MoE MLP = 16,384 Γ— (36+1) Γ— (8192 + 3Γ—2048) Γ— 2 = 16.19 GB +``` + +##### Full Transformer Layer (without recompute) + +For a MoE layer, the total is the sum of all components: + +``` +layer_act = 3Γ—LayerNorm + Attention + 2Γ—ResidualAdd + Router + MoE_MLP + = 3Γ—sbh + attn_act + 2Γ—sbh + sbh + moe_mlp_act + = 6Γ—sbh + attn_act + moe_mlp_act +``` + +For a dense layer: same but with Dense MLP instead of Router + MoE MLP. + +##### Full Layer Activation Summary + +| Component | Formula | Typical Size (MoE 4.5T) | +|-----------|---------|------------------------| +| LayerNorm (Γ—3) | `3 Γ— sbh` | 0.75 GB | +| Residual Add (Γ—2) | `2 Γ— sbh` | 0.50 GB | +| Router | `sbh` | 0.25 GB | +| Attention (Flash, MQA) | `tokens Γ— (Q+2KV+O+softmax) Γ— 2` | 0.51 GB | +| MoE MLP (SwiGLU) | `tokens Γ— (topk+1) Γ— (H+3Γ—FFN_e) Γ— 2` | 16.19 GB | +| **Full MoE layer** | **sum** | **18.20 GB** | +| With full recompute | `sbh` (checkpoint only) | 0.25 GB | + + +### Pipeline Schedule Memory Scaling + +With pipeline parallelism, multiple microbatches are in-flight simultaneously, each requiring stored activations. + +#### 1F1B Schedule + +In the 1F1B (one-forward-one-backward) schedule, the first pipeline stage (rank 0) accumulates `PP` microbatches during the warmup phase before starting any backward passes. This means peak activation memory requires storing activations for `PP` microbatches. + +``` +base_activation = sum of per-layer activations across all layers on this rank +peak_activation = base_activation Γ— PP +``` + +#### VPP (Virtual Pipeline Parallelism) Overhead + +With interleaved scheduling (VPP > 1), there is a small memory overhead because more microbatches can be partially in-flight: + +``` +interleaved_penalty = 1 + (PP - 1) / (PP Γ— VPP) +``` + +For VPP=1: penalty = 1 + (PP-1)/PP (significant overhead) +For VPP=20: penalty = 1 + (PP-1)/(PPΓ—20) β‰ˆ 1.04 (nearly negligible) + +#### Gradient Accumulation Saving + +When gradient accumulation (GA) steps are fewer than PP stages, the pipeline isn't fully filled: + +```python +GA = GBS / (MBS Γ— DP) +ga_saving = 1.0 if GA >= PP else GA / PP +``` + +#### Final Activation Memory Formula + +``` +total_activation = base_activation Γ— PP Γ— interleaved_penalty Γ— ga_saving +``` + + +### Recomputation Support + +Activation recomputation trades compute for memory by discarding intermediate activations during forward and recomputing them during backward. + +#### Full Recompute (`recompute_granularity="full"`) + +When a layer is fully recomputed, only the **input tensor** is stored as a checkpoint: + +``` +recomputed_layer_act = sbh = MBS Γ— S/(TPΓ—CP) Γ— H Γ— 2 bytes +``` + +This is dramatically smaller than the full activation. For example, an MoE layer drops from ~18 GB to 0.25 GB. + +#### Partial Recompute + +The `recompute_num_layers` setting controls how many layers per VPP stage are recomputed: + +```python +for each layer on this rank: + local_idx = layer_index % layers_per_vpp_stage + if recompute_granularity == "full" and local_idx < recompute_num_layers: + act += input_act_per_layer # just sbh (checkpoint) + else: + act += full_layer_act # all intermediates +``` + +#### With Recompute: Total Memory + +``` +total_with_recompute = (L/PP Γ— sbh) Γ— PP Γ— interleaved_penalty Γ— ga_saving + + recompute_working_memory (1 layer's full activation, temporary) + + embedding_act (stage 0 only) +``` + +The recompute working memory is transient β€” only one layer's full intermediates exist at a time during backward. + +### Memory Formulas Reference + +Summary of all memory components for one GPU: + +``` +β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” +β”‚ Total GPU Memory β”‚ +β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ +β”‚ β”‚ +β”‚ 1. Parameters (BF16) β”‚ +β”‚ = params_on_rank Γ— 2 bytes β”‚ +β”‚ β”‚ +β”‚ 2. Gradients (BF16) β”‚ +β”‚ = params_on_rank Γ— 2 bytes β”‚ +β”‚ β”‚ +β”‚ 3. Optimizer States (FP32, sharded across DP) β”‚ +β”‚ = params_on_rank Γ— 10 / DP bytes β”‚ +β”‚ (master weights: 2B + Adam m: 4B + Adam v: 4B) β”‚ +β”‚ β”‚ +β”‚ 4. Activations β”‚ +β”‚ = Ξ£(per-layer act) Γ— PP Γ— VPP_penalty Γ— GA_saving β”‚ +β”‚ + embedding/output activations (stage-dependent) β”‚ +β”‚ β”‚ +β”‚ 5. Transient buffers (not in projection) β”‚ +β”‚ - A2A dispatch/combine buffers β”‚ +β”‚ - Communication scratch space β”‚ +β”‚ - PyTorch allocator fragmentation β”‚ +β”‚ β”‚ +β”‚ Total = (1) + (2) + (3) + (4) β”‚ +β”‚ Reported as: Param+Optimizer + Activation = Projected Total β”‚ +β”‚ β”‚ +β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ +``` + +--- + +## Performance Projection + + +### Overview + +The performance projection tool: + +1. **Benchmarks** transformer layers on a single node to measure forward/backward pass times +2. **Simulates** pipeline parallelism scheduling (including zero-bubble optimization) +3. **Projects** performance to multi-node configurations by modeling: + - Data Parallelism (DP) scaling + - Gradient AllReduce communication overhead + - Expert Parallelism (EP) All-to-All communication overhead + - Inter-node vs intra-node communication differences + +This allows you to estimate training performance on larger clusters without actually running on them. + +### How It Works + +#### 1. Configuration Reduction If the parallelism configuration requires multiple nodes (e.g., PP=3 needs 3 nodes), the tool automatically reduces the config for single-node benchmarking: - **Pipeline Parallelism (PP)**: Reduced to fit on 1 node, PP overhead estimated analytically - **Expert Parallelism (EP)**: Reduced if needed, All-to-All overhead added back -### 2. Layer Benchmarking +#### 2. Layer Benchmarking The tool benchmarks each transformer layer type: - Dense attention layers - MoE (Mixture of Experts) layers - Measures forward and backward pass times separately +- Also benchmarks embedding and output layers -### 3. Pipeline Simulation +#### 3. Pipeline Simulation For PP > 1, the tool simulates the pipeline schedule to account for: - Pipeline bubble overhead - Microbatch interleaving - Zero-bubble scheduling (if enabled) -### 4. Data Parallel Scaling +#### 4. Data Parallel Scaling The projection models how performance scales with additional nodes: @@ -106,11 +497,9 @@ The projection models how performance scales with additional nodes: Projected Time = (Base Time / DP_scaling_factor) + Communication Overheads ``` -## Scaling Mechanisms +### Scaling Mechanisms -The tool models the following parallelism dimensions and their communication patterns: - -### Tensor Parallelism (TP) +#### Tensor Parallelism (TP) **What it does**: Splits individual layer weights across GPUs within a node. @@ -118,7 +507,7 @@ The tool models the following parallelism dimensions and their communication pat **Communication**: AllReduce within TP group (typically intra-node, fast). -### Pipeline Parallelism (PP) +#### Pipeline Parallelism (PP) **What it does**: Distributes layers across pipeline stages. Each stage processes microbatches in sequence. @@ -128,7 +517,7 @@ The tool models the following parallelism dimensions and their communication pat - Simulates the actual 1F1B or zero-bubble schedule with proper send/receive synchronization - Accounts for pipeline bubble overhead and microbatch interleaving -### Expert Parallelism (EP) +#### Expert Parallelism (EP) **What it does**: Distributes MoE experts across GPUs. Each GPU holds a subset of experts. @@ -143,7 +532,7 @@ The tool models the following parallelism dimensions and their communication pat All-to-All Message Size = tokens Γ— hidden_size Γ— top_k Γ— 2 (BF16) ``` -### Data Parallelism (DP) +#### Data Parallelism (DP) **What it does**: Replicates the model across DP groups. Each group processes different data batches. @@ -159,13 +548,13 @@ Gradient AllReduce Size = num_params Γ— 4 (FP32 gradients) **Optimization**: If `overlap_grad_reduce=True` (default), gradient AllReduce is overlapped with backward computation and not on the critical path. -### Context Parallelism (CP) +#### Context Parallelism (CP) **What it does**: Splits sequence length across GPUs for long-context training. **How it's modeled**: CP affects the GPU topology for communication routing. Currently included in minimum GPU requirements calculation. -## Communication Modeling +### Communication Modeling The tool uses analytical models to estimate collective communication times: @@ -176,33 +565,177 @@ The tool uses analytical models to estimate collective communication times: | P2P Send/Recv | PP (activations) | Point-to-point latency + bandwidth | Communication times differ significantly based on: -- **Intra-node**: Fast (e.g., NVLink, xGMI) -- **Inter-node**: Slower (e.g., InfiniBand, RoCE) +- **Intra-node**: Fast (e.g., NVLink, UALink, xGMI) +- **Inter-node**: Slower (e.g., InfiniBand, RoCE) + +Custom hardware parameters can be provided via `--hardware-config `. -### Key Concepts +#### Key Concepts -#### Minimum Nodes Required +##### Minimum Nodes Required -The minimum nodes required is determined by: ``` Min Nodes = ceil(TP Γ— PP Γ— EP Γ— CP / GPUs_per_node) ``` -#### Scaling Behavior +##### Scaling Behavior - **DP scaling**: Linear speedup. Doubling DP halves iteration time (minus communication overhead). - **PP scaling**: Happens in multiples of pipeline replicas. With PP=3, you need 3, 6, 9... nodes to increase scaling. -- **EP scaling**: Divides the experts on EP nodes. +- **EP scaling**: Divides the experts across EP nodes. + +### Pipeline Schedule Simulator + +The pipeline simulator (`simulator.py`) simulates the execution of pipeline parallelism schedules to calculate step time and bubble ratio. + +#### Schedule Algorithms + +| Algorithm | Description | Use Case | +|-----------|-------------|----------| +| **1F1B** | Standard one-forward-one-backward | Default pipeline schedule | +| **Interleaved 1F1B** | Multiple chunks per rank (VPP > 1) | Reduced bubble ratio | +| **Zero-Bubble** | Splits backward into B + W | Minimal bubble overhead | + +#### Zero-Bubble Scheduling + +Zero-bubble minimizes pipeline bubbles by separating the backward pass: +- **B (Input Gradient):** Compute gradients w.r.t. input activations +- **W (Weight Gradient):** Compute gradients w.r.t. weights + +This allows more flexible scheduling because W doesn't depend on receiving gradients from the next stage. By default, backward time is split 50/50 between B and W. + +Two implementations are available: +1. **Simple Zero-Bubble Simulator** β€” basic F-B-W pattern with warmup/steady/cooldown phases +2. **Megatron ILP-Based Scheduler** β€” graph-based schedule optimization with memory-aware scheduling using Megatron's actual zero-bubble scheduler + +#### P2P Communication in Pipeline Simulation + +The pipeline simulator uses a **fixed small constant** (~0.1 ms) for P2P communication, NOT the analytical `sendrecv` model. This is because: +1. P2P communication is typically **overlapped with compute** in optimized schedules +2. The simulator focuses on **schedule ordering and bubble calculation** +3. P2P time is **small relative to F/B/W times** for large models + +However, when the benchmark PP differs from the target PP (e.g., benchmark PP=1, target PP=6), the **analytical `sendrecv` model** is used to estimate the PP communication overhead that was not captured in the benchmark: + +``` +PP overhead = additional_stages Γ— 2 (fwd+bwd) Γ— microbatches Γ— sendrecv(activation_size) +``` + +P2P communication becomes significant when PP stages span nodes (inter-node P2P has 10-100Γ— higher latency than intra-node). + +### Overall Performance Prediction Flow + +``` +1. Load Configuration + └── Parse YAML config, extract parallelism settings + +2. Single-Node Benchmarking + β”œβ”€β”€ If config requires multiple nodes: + β”‚ └── Reduce PP to 1, possibly rescale EP to fit on 1 node + β”œβ”€β”€ Limit layers (1 dense + 1 MoE for efficiency) + └── Benchmark forward + backward times + +3. Extrapolate to Full Model + └── Multiply per-layer times by total layer count + +4. Pipeline Schedule Simulation (if PP > 1) + β”œβ”€β”€ Build chunk time matrix (per-rank, per-chunk) + β”œβ”€β”€ Select scheduler (1F1B, Interleaved, Zero-Bubble) + └── Simulate to get step_time_ms and bubble_ratio + +5. Add Communication Overhead (if config was reduced) + β”œβ”€β”€ PP overhead: P2P communication between stages + └── EP overhead: Additional All-to-All for larger EP + +6. Multinode Scaling Projection + β”œβ”€β”€ Calculate DP scaling factor + β”œβ”€β”€ Scale compute time: projected = base Γ— (min_dp / target_dp) + β”œβ”€β”€ Add gradient AllReduce (if not overlapped) + └── Report projected iteration time and tokens/s +``` + +#### Performance Formula + +``` +Projected_Time = Base_Time Γ— (Min_DP / Target_DP) + Communication_Overhead + +Where: + Base_Time = Pipeline simulation time (includes PP bubbles) + Min_DP = DP at minimum node configuration + Target_DP = DP at target node configuration + Communication_Overhead = Gradient AllReduce (if not overlapped) + + MoE All-to-All overhead (if EP spans nodes) +``` + +#### Example Calculation + +**Configuration:** DeepSeek V2 Lite β€” TP=1, PP=3, EP=8, CP=1 β€” GBS=640, MBS=4, Seq=4096 + +``` +Step 1: Minimum Nodes + GPUs required = 1 Γ— 3 Γ— 8 Γ— 1 = 24 GPUs = 3 nodes + Min DP = 24 / (1 Γ— 3 Γ— 1) = 8 + +Step 2: Target Configuration (6 nodes) + Total GPUs = 48 + Target DP = 48 / (1 Γ— 3 Γ— 1) = 16 + DP Scaling = 16 / 8 = 2Γ— + +Step 3: Projected Time + Base_Time (from pipeline simulation) = 5500 ms + Projected_Time = 5500 Γ— (8 / 16) = 2750 ms + Tokens/s = (640 Γ— 4096) / 2.750 = 953,018 tokens/s +``` + +--- ## Example Output +### Memory Projection + ``` ==================================================================================================== -[Primus:Performance Projection] Configuration Summary: - Benchmark Config: PP=1, EP=8, TP=1, CP=1, DP=1 (1 node) - Target Config: PP=1, EP=8, TP=1, CP=1, DP=4 (4 nodes) - Benchmark Microbatches: 160 (global_batch=640, micro_batch=4, benchmark_dp=1) +[Primus:Projection] Component-wise Profiling Results (Rank 0): +==================================================================================================== + + Total Number of Parameters: 15.654321 Billion (15,654,321,024) + + [embedding] + Params: 0.819200 Billion (819,200,000) + Activation Memory: 0.2500 GB + + [dense_transformer_layer] + Params: 0.302000 Billion (302,000,000) + Activation Memory: 2.1250 GB + + [layer_norm] + Params: 0.000016 Billion (16,384) + Activation Memory: 0.2500 GB + + [self_attention] + Params: 0.134218 Billion (134,217,728) + Activation Memory: 0.5100 GB + + [mlp] + Params: 0.167772 Billion (167,772,160) + Activation Memory: 0.8650 GB + + [moe_transformer_layer] + Params: 1.001400 Billion (1,001,400,000) + Activation Memory: 18.2000 GB + +==================================================================================================== +[Primus:Projection] Memory Projection Summary on Rank 0: + Params: 20.850000 Billion (20,850,000,000) + Param+Optimizer Memory: 83.7400 GB + Activation Memory (per batch size 4, seq len 16384): 36.7500 GB + Projected Total Memory: 120.4900 GB +==================================================================================================== +``` +### Performance Projection + +``` ==================================================================================================== Multinode Scaling Projection Results ==================================================================================================== @@ -226,12 +759,36 @@ Multinode Scaling Projection Results ==================================================================================================== ``` +--- + +## Assumptions and Limitations + +### Assumptions + +1. **Linear DP Scaling** β€” Compute time scales linearly with 1/DP (ideal weak scaling) +2. **Communication Model** β€” Bandwidth efficiency is constant (default 91%); inter-node communication uses switch topology; all NICs are used in parallel for inter-node traffic +3. **Pipeline Bubbles** β€” B/W split is 50/50 for zero-bubble scheduling; P2P communication time is small relative to compute +4. **Gradient AllReduce** β€” By default overlapped with compute (`overlap_grad_reduce=True`); if disabled, added to critical path +5. **MoE All-to-All** β€” All-to-All time scales with EP size; per-peer latency overhead is constant + +### Limitations + +1. **Single-Node Benchmark Accuracy** β€” Benchmarking with reduced PP/EP may not capture all behaviors; layer timing variance is assumed uniform +2. **Communication Contention** β€” Model doesn't account for network contention; assumes dedicated bandwidth per collective +3. **Memory Pressure** β€” Memory impact on performance not fully modeled; activation recomputation overhead not considered in performance +4. **Hardware Heterogeneity** β€” Assumes homogeneous nodes; GPU frequency variations not modeled + +--- + ## Tips +- **Start with memory projection**: Run `projection memory` first to verify your model fits in GPU memory before benchmarking. - **Start with 1 node**: Always benchmark on 1 node first to establish baseline performance. -- **Understand scaling limits**: DP scaling is limited by global_batch_size / micro_batch_size. If you run out of microbatches, adding more nodes won't help. -- **Check minimum nodes**: If your config requires multiple nodes (e.g., PP=4 with 8 GPUs/node), projection will automatically reduce PP for benchmarking. +- **Understand scaling limits**: DP scaling is limited by `global_batch_size / micro_batch_size`. If you run out of microbatches, adding more nodes won't help. +- **Check minimum nodes**: If your config requires multiple nodes (e.g., PP=4 with 8 GPUs/node), the performance projection will automatically reduce PP for benchmarking. - **Pipeline scaling**: With PP > 1, you can only scale in multiples of the pipeline replica size. +- **Recomputation trade-off**: Full recompute dramatically reduces activation memory (e.g., 18 GB β†’ 0.25 GB per MoE layer) at the cost of ~33% more compute. +- **MoE activation dominance**: For MoE models, the MoE MLP activation (scaled by `topk`) typically dominates the per-layer activation budget. Consider recomputation if memory is tight. ## Related Documentation From d32360ee1cd4710cee7325bd3f41d81aa932b1fd Mon Sep 17 00:00:00 2001 From: vidushi8 Date: Thu, 5 Mar 2026 19:51:04 -0600 Subject: [PATCH 09/11] update yamls to fix regressions and standardize --- .../configs/MI300X/deepseek_v2_lite-BF16-pretrain.yaml | 6 +++++- .../megatron/configs/MI355X/deepseek_v2-BF16-pretrain.yaml | 4 ++++ .../configs/MI355X/deepseek_v2_lite-BF16-pretrain.yaml | 7 +++++-- .../megatron/configs/MI355X/deepseek_v3-BF16-pretrain.yaml | 5 +++++ .../configs/MI355X/llama3.1_70B-BF16-pretrain.yaml | 4 ++-- .../configs/MI300X/llama3.1_70B-BF16-pretrain.yaml | 1 + .../configs/MI300X/llama3.1_70B-FP8-pretrain.yaml | 1 + .../configs/MI300X/llama3.1_8B-BF16-pretrain.yaml | 1 + .../configs/MI300X/llama3.1_8B-FP8-pretrain.yaml | 1 + .../configs/MI355X/llama3.1_70B-BF16-pretrain.yaml | 3 ++- .../configs/MI355X/llama3.1_70B-FP8-pretrain.yaml | 3 ++- .../configs/MI355X/llama3.1_8B-BF16-pretrain.yaml | 1 + .../configs/MI355X/llama3.1_8B-FP8-pretrain.yaml | 1 + 13 files changed, 31 insertions(+), 7 deletions(-) diff --git a/examples/megatron/configs/MI300X/deepseek_v2_lite-BF16-pretrain.yaml b/examples/megatron/configs/MI300X/deepseek_v2_lite-BF16-pretrain.yaml index d97c22d1e..786cb082d 100644 --- a/examples/megatron/configs/MI300X/deepseek_v2_lite-BF16-pretrain.yaml +++ b/examples/megatron/configs/MI300X/deepseek_v2_lite-BF16-pretrain.yaml @@ -88,7 +88,11 @@ modules: # stage 3 is completely no gpu-cpu sync in MoE, but cost more memory # stage 2 is recommended for better performance turbo_sync_free_moe_stage: 2 - + + # remove once super flag is functional again + moe_use_fused_router_with_aux_score: true + moe_permute_fusion: true + # Cross entropy flags cross_entropy_fusion_impl: "te" cross_entropy_loss_fusion: true diff --git a/examples/megatron/configs/MI355X/deepseek_v2-BF16-pretrain.yaml b/examples/megatron/configs/MI355X/deepseek_v2-BF16-pretrain.yaml index 5f7d1299c..f2b431c0d 100644 --- a/examples/megatron/configs/MI355X/deepseek_v2-BF16-pretrain.yaml +++ b/examples/megatron/configs/MI355X/deepseek_v2-BF16-pretrain.yaml @@ -93,6 +93,10 @@ modules: # stage 2 is recommended for better performance turbo_sync_free_moe_stage: 1 + # remove once super flag is functional again + moe_use_fused_router_with_aux_score: true + moe_permute_fusion: true + # Cross entropy flags cross_entropy_fusion_impl: "te" cross_entropy_loss_fusion: true diff --git a/examples/megatron/configs/MI355X/deepseek_v2_lite-BF16-pretrain.yaml b/examples/megatron/configs/MI355X/deepseek_v2_lite-BF16-pretrain.yaml index 54051c9c7..775804e48 100644 --- a/examples/megatron/configs/MI355X/deepseek_v2_lite-BF16-pretrain.yaml +++ b/examples/megatron/configs/MI355X/deepseek_v2_lite-BF16-pretrain.yaml @@ -52,7 +52,7 @@ modules: valid_data_path: null test_data_path: null - moe_use_legacy_grouped_gemm: false + moe_use_legacy_grouped_gemm: true # MLA multi_latent_attention: true @@ -73,7 +73,7 @@ modules: # Turbo enable_primus_turbo: true use_turbo_attention: false - use_turbo_grouped_mlp: true + use_turbo_grouped_mlp: false # deepep use_turbo_deepep: true @@ -87,7 +87,10 @@ modules: # sync-free moe support stage 1-2, 0 means not use sync-free moe # stage 2 is recommended for better performance turbo_sync_free_moe_stage: 1 + + # remove once super flag is functional again moe_use_fused_router_with_aux_score: true + moe_permute_fusion: true # Cross entropy flags cross_entropy_fusion_impl: "te" diff --git a/examples/megatron/configs/MI355X/deepseek_v3-BF16-pretrain.yaml b/examples/megatron/configs/MI355X/deepseek_v3-BF16-pretrain.yaml index 61bd8aa5d..7e863b2ae 100644 --- a/examples/megatron/configs/MI355X/deepseek_v3-BF16-pretrain.yaml +++ b/examples/megatron/configs/MI355X/deepseek_v3-BF16-pretrain.yaml @@ -84,6 +84,11 @@ modules: # stage 2 is recommended for better performance turbo_sync_free_moe_stage: 1 + # remove once super flag is functional again + moe_use_fused_router_with_aux_score: true + moe_permute_fusion: true + + # Cross entropy flags cross_entropy_fusion_impl: "te" cross_entropy_loss_fusion: true diff --git a/examples/megatron/configs/MI355X/llama3.1_70B-BF16-pretrain.yaml b/examples/megatron/configs/MI355X/llama3.1_70B-BF16-pretrain.yaml index c707572da..f25bf4eb9 100644 --- a/examples/megatron/configs/MI355X/llama3.1_70B-BF16-pretrain.yaml +++ b/examples/megatron/configs/MI355X/llama3.1_70B-BF16-pretrain.yaml @@ -17,8 +17,8 @@ modules: log_avg_reset_interval: 50 train_iters: 50 - micro_batch_size: 7 - global_batch_size: 56 + micro_batch_size: 8 + global_batch_size: 128 seq_length: 8192 max_position_embeddings: 8192 diff --git a/examples/torchtitan/configs/MI300X/llama3.1_70B-BF16-pretrain.yaml b/examples/torchtitan/configs/MI300X/llama3.1_70B-BF16-pretrain.yaml index 21970087a..d082dc5f7 100644 --- a/examples/torchtitan/configs/MI300X/llama3.1_70B-BF16-pretrain.yaml +++ b/examples/torchtitan/configs/MI300X/llama3.1_70B-BF16-pretrain.yaml @@ -27,6 +27,7 @@ modules: seq_len: 8192 mock_data: false steps: 50 + gc_freq: 100 optimizer: lr: 1.5e-4 diff --git a/examples/torchtitan/configs/MI300X/llama3.1_70B-FP8-pretrain.yaml b/examples/torchtitan/configs/MI300X/llama3.1_70B-FP8-pretrain.yaml index 6e09f2dcf..0ca09d214 100644 --- a/examples/torchtitan/configs/MI300X/llama3.1_70B-FP8-pretrain.yaml +++ b/examples/torchtitan/configs/MI300X/llama3.1_70B-FP8-pretrain.yaml @@ -27,6 +27,7 @@ modules: seq_len: 8192 mock_data: false steps: 50 + gc_freq: 100 optimizer: lr: 1.5e-4 diff --git a/examples/torchtitan/configs/MI300X/llama3.1_8B-BF16-pretrain.yaml b/examples/torchtitan/configs/MI300X/llama3.1_8B-BF16-pretrain.yaml index 927f96662..43a7e1fcb 100644 --- a/examples/torchtitan/configs/MI300X/llama3.1_8B-BF16-pretrain.yaml +++ b/examples/torchtitan/configs/MI300X/llama3.1_8B-BF16-pretrain.yaml @@ -28,6 +28,7 @@ modules: seq_len: 8192 mock_data: false steps: 50 + gc_freq: 100 activation_checkpoint: mode: "none" # ["none", "selective", "full"] diff --git a/examples/torchtitan/configs/MI300X/llama3.1_8B-FP8-pretrain.yaml b/examples/torchtitan/configs/MI300X/llama3.1_8B-FP8-pretrain.yaml index 35dfe6ade..29c6dd346 100644 --- a/examples/torchtitan/configs/MI300X/llama3.1_8B-FP8-pretrain.yaml +++ b/examples/torchtitan/configs/MI300X/llama3.1_8B-FP8-pretrain.yaml @@ -27,6 +27,7 @@ modules: seq_len: 8192 mock_data: false steps: 50 + gc_freq: 100 activation_checkpoint: mode: "none" # ["none", "selective", "full"] diff --git a/examples/torchtitan/configs/MI355X/llama3.1_70B-BF16-pretrain.yaml b/examples/torchtitan/configs/MI355X/llama3.1_70B-BF16-pretrain.yaml index 1e4465d38..9b5ea8f3a 100644 --- a/examples/torchtitan/configs/MI355X/llama3.1_70B-BF16-pretrain.yaml +++ b/examples/torchtitan/configs/MI355X/llama3.1_70B-BF16-pretrain.yaml @@ -26,7 +26,8 @@ modules: local_batch_size: 3 seq_len: 8192 mock_data: false - steps: 10 + steps: 50 + gc_freq: 100 profiling: enable_profiling: false diff --git a/examples/torchtitan/configs/MI355X/llama3.1_70B-FP8-pretrain.yaml b/examples/torchtitan/configs/MI355X/llama3.1_70B-FP8-pretrain.yaml index 954b3826d..432c21b1d 100644 --- a/examples/torchtitan/configs/MI355X/llama3.1_70B-FP8-pretrain.yaml +++ b/examples/torchtitan/configs/MI355X/llama3.1_70B-FP8-pretrain.yaml @@ -32,7 +32,8 @@ modules: global_batch_size: 96 seq_len: 8192 mock_data: false - steps: 10 + steps: 50 + gc_freq: 100 optimizer: lr: 1.5e-4 diff --git a/examples/torchtitan/configs/MI355X/llama3.1_8B-BF16-pretrain.yaml b/examples/torchtitan/configs/MI355X/llama3.1_8B-BF16-pretrain.yaml index fd9993ddf..50b276070 100644 --- a/examples/torchtitan/configs/MI355X/llama3.1_8B-BF16-pretrain.yaml +++ b/examples/torchtitan/configs/MI355X/llama3.1_8B-BF16-pretrain.yaml @@ -28,6 +28,7 @@ modules: seq_len: 8192 mock_data: false steps: 50 + gc_freq: 100 activation_checkpoint: mode: "none" # ["none", "selective", "full"] diff --git a/examples/torchtitan/configs/MI355X/llama3.1_8B-FP8-pretrain.yaml b/examples/torchtitan/configs/MI355X/llama3.1_8B-FP8-pretrain.yaml index 6e7202701..aa71f9c83 100644 --- a/examples/torchtitan/configs/MI355X/llama3.1_8B-FP8-pretrain.yaml +++ b/examples/torchtitan/configs/MI355X/llama3.1_8B-FP8-pretrain.yaml @@ -27,6 +27,7 @@ modules: seq_len: 8192 mock_data: false steps: 50 + gc_freq: 100 activation_checkpoint: mode: "none" # ["none", "selective", "full"] From bc4c861e33f0145fa9015d307aa0c7f11e7abeec Mon Sep 17 00:00:00 2001 From: WangLingxun Date: Fri, 6 Mar 2026 16:13:49 +0800 Subject: [PATCH 10/11] fix(megatron): patch validate_args and add ROCM argument validation (#581) Hook Megatron validate_args alongside parse_args so Primus-injected arguments are validated consistently, and run additional ROCM-specific argument checks during initialization. --- .../megatron/megatron_base_trainer.py | 24 +++++++++++++++++-- 1 file changed, 22 insertions(+), 2 deletions(-) diff --git a/primus/backends/megatron/megatron_base_trainer.py b/primus/backends/megatron/megatron_base_trainer.py index 35909d120..26129c910 100644 --- a/primus/backends/megatron/megatron_base_trainer.py +++ b/primus/backends/megatron/megatron_base_trainer.py @@ -23,17 +23,37 @@ def init(self): # log_dict_aligned("Backend arguments", self.backend_args) def _patch_parse_args(self): - """Patch Megatron's parse_args to return pre-configured Primus arguments.""" + """ + This function patches Megatron's parse_args to return pre-configured Primus arguments. + It also validates the arguments on ROCM. + """ import megatron.training.arguments as megatron_args # type: ignore import megatron.training.initialize as megatron_init # type: ignore + from primus.modules.trainer.megatron.utils import validate_args_on_rocm + log_rank_0("Patching Megatron-LM parse_args()") patched_parse_args = lambda *args, **kwargs: ( log_rank_0("parse_args() called; returning Primus arguments") or self.backend_args ) + original_validate_args = megatron_args.validate_args + + def patched_validate_args(*args, **kwargs): + validated_args = original_validate_args(*args, **kwargs) + parsed_args = args[0] if args else kwargs.get("args", None) + if parsed_args is not None: + log_rank_0("validate_args() called; validating on ROCM") + validate_args_on_rocm(parsed_args) + return validated_args + megatron_args.parse_args = patched_parse_args megatron_init.parse_args = patched_parse_args - log_rank_0(f"Patched parse_args(); Primus provided {len(vars(self.backend_args))} arguments") + megatron_args.validate_args = patched_validate_args + megatron_init.validate_args = patched_validate_args + + log_rank_0( + f"Patched parse_args()/validate_args(); Primus provided {len(vars(self.backend_args))} arguments" + ) From ecd27b031ff9d5cc1355a865170d93cb48e5484a Mon Sep 17 00:00:00 2001 From: clairesonglee Date: Tue, 10 Mar 2026 04:41:23 +0000 Subject: [PATCH 11/11] Add unit tests for mamba and hybrid models --- tests/trainer/test_megatron_trainer.py | 28 ++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/tests/trainer/test_megatron_trainer.py b/tests/trainer/test_megatron_trainer.py index ad3acc280..a3af36d82 100644 --- a/tests/trainer/test_megatron_trainer.py +++ b/tests/trainer/test_megatron_trainer.py @@ -324,6 +324,34 @@ def test_turbo_deepep(self): ], ) + def test_zebra_llama_1B(self): + run_script( + self.__class__.__name__, + "zebra_llama_1B", + exp_path="examples/megatron/configs/MI300X/zebra_llama_1B-pretrain.yaml", + env_override={}, + extra_args=[ + "--num_layers", + "4", + "--train_iters", + "3", + ], + ) + + def test_mamba_370M(self): + run_script( + self.__class__.__name__, + "mamba_370M", + exp_path="examples/megatron/configs/MI300X/mamba_370M-pretrain.yaml", + env_override={}, + extra_args=[ + "--num_layers", + "4", + "--train_iters", + "3", + ], + ) + class TestMegatronTrainerDeterministic(PrimusUT): def __init__(self, *args, **kwargs):