fix(recipe): prevent Qwen3.5-VL FSDP SFT OOM (#5508)

Signed-off-by: Chen Cui <chcui@nvidia.com>
Signed-off-by: NeMo Bot <nemo-bot@nvidia.com>

Author: cuichenx

src/megatron/bridge/recipes/qwen_vl/h100/qwen35_vl.py

@@ -951,7 +951,9 @@ def qwen35_vl_35b_a3b_sft_2gpu_h100_bf16_fsdp_config() -> ConfigContainer:
     cfg.model.attention_backend = "auto"
     cfg.model.gradient_accumulation_fusion = True
     cfg.model.cross_entropy_loss_fusion = True
-    cfg.model.cross_entropy_fusion_impl = "native"
+    # Native fused cross entropy materializes a full FP32 vocabulary-sized
+    # softmax buffer. Use the TE kernel to keep the 248K-token loss bounded.
+    cfg.model.cross_entropy_fusion_impl = "te"
 
     # MoE settings
     cfg.model.moe_token_dispatcher_type = "alltoall"
@@ -965,10 +967,12 @@ def qwen35_vl_35b_a3b_sft_2gpu_h100_bf16_fsdp_config() -> ConfigContainer:
     cfg.model.moe_router_padding_for_fp8 = False
 
     # Memory saving
-    cfg.model.recompute_granularity = None
+    # Full SFT retains FP32 optimizer state alongside the FSDP communication
+    # pools, so recompute each transformer layer to bound live activations.
+    cfg.model.recompute_granularity = "full"
     cfg.model.recompute_modules = None
-    cfg.model.recompute_method = None
-    cfg.model.recompute_num_layers = None
+    cfg.model.recompute_method = "uniform"
+    cfg.model.recompute_num_layers = 1
     cfg.model.fine_grained_activation_offloading = False
     cfg.model.offload_modules = None
 

tests/unit_tests/recipes/qwen_vl/test_qwen35_vl_recipes.py

@@ -431,6 +431,12 @@ def test_qwen35_vl_35b_a3b_fsdp_sft_defaults(monkeypatch: pytest.MonkeyPatch):
     assert cfg.model.sequence_parallel is False
     assert cfg.model.moe_token_dispatcher_type == "alltoall"
     assert cfg.model.moe_router_fusion is True
+    assert cfg.model.cross_entropy_loss_fusion is True
+    assert cfg.model.cross_entropy_fusion_impl == "te"
+    assert cfg.model.recompute_granularity == "full"
+    assert cfg.model.recompute_modules is None
+    assert cfg.model.recompute_method == "uniform"
+    assert cfg.model.recompute_num_layers == 1
     assert cfg.ddp.use_megatron_fsdp is True
     assert cfg.ddp.fsdp_double_buffer is True
     assert cfg.ddp.megatron_fsdp_max_pool_double_buffer is True