File tree Expand file tree Collapse file tree
src/megatron/bridge/recipes/qwen_vl/h100
tests/unit_tests/recipes/qwen_vl Expand file tree Collapse file tree Original file line number Diff line number Diff line change @@ -1190,7 +1190,9 @@ def qwen35_vl_35b_a3b_sft_2gpu_h100_bf16_fsdp_config() -> ConfigContainer:
11901190 cfg .model .attention_backend = "auto"
11911191 cfg .model .gradient_accumulation_fusion = True
11921192 cfg .model .cross_entropy_loss_fusion = True
1193- cfg .model .cross_entropy_fusion_impl = "native"
1193+ # Native fused cross entropy materializes a full FP32 vocabulary-sized
1194+ # softmax buffer. Use the TE kernel to keep the 248K-token loss bounded.
1195+ cfg .model .cross_entropy_fusion_impl = "te"
11941196
11951197 # MoE settings
11961198 cfg .model .moe_token_dispatcher_type = "alltoall"
@@ -1204,10 +1206,12 @@ def qwen35_vl_35b_a3b_sft_2gpu_h100_bf16_fsdp_config() -> ConfigContainer:
12041206 cfg .model .moe_router_padding_for_fp8 = False
12051207
12061208 # Memory saving
1207- cfg .model .recompute_granularity = None
1209+ # Full SFT retains FP32 optimizer state alongside the FSDP communication
1210+ # pools, so recompute each transformer layer to bound live activations.
1211+ cfg .model .recompute_granularity = "full"
12081212 cfg .model .recompute_modules = None
1209- cfg .model .recompute_method = None
1210- cfg .model .recompute_num_layers = None
1213+ cfg .model .recompute_method = "uniform"
1214+ cfg .model .recompute_num_layers = 1
12111215 cfg .model .fine_grained_activation_offloading = False
12121216 cfg .model .offload_modules = None
12131217
Original file line number Diff line number Diff line change @@ -561,6 +561,12 @@ def test_qwen35_vl_35b_a3b_fsdp_sft_defaults(monkeypatch: pytest.MonkeyPatch):
561561 assert cfg .model .sequence_parallel is False
562562 assert cfg .model .moe_token_dispatcher_type == "alltoall"
563563 assert cfg .model .moe_router_fusion is True
564+ assert cfg .model .cross_entropy_loss_fusion is True
565+ assert cfg .model .cross_entropy_fusion_impl == "te"
566+ assert cfg .model .recompute_granularity == "full"
567+ assert cfg .model .recompute_modules is None
568+ assert cfg .model .recompute_method == "uniform"
569+ assert cfg .model .recompute_num_layers == 1
564570 assert cfg .ddp .use_megatron_fsdp is True
565571 assert cfg .ddp .fsdp_double_buffer is True
566572 assert cfg .ddp .megatron_fsdp_max_pool_double_buffer is True
You can’t perform that action at this time.
0 commit comments