Recipe Optimization Request: GPT-OSS 20B DAPT on 8x B200

[bolivasNV]

Hi!
I am working on a pre-training job for GPT-OSS 20B (25B tokens, trained in 3 phases). I need your feedback to optimize the pre-training recipe to reduce our current 2-week timeline on 8x B200 GPUs. We are currently running at a GBS of 16, which results in a 2-week timeline. Any setting higher than GBS 256 triggers an immediate OOM error. Also, if we increase the batch size to our target of 256, the projected training time actually augments to 3 weeks.

Current Recipe Hyperparameters:

• Parallelism: TP=1, PP=1, CP=1, EP=2, DP=4.
• Batching: GBS=16, MBS=1, Seq Length=4096.
• Precision: bf16-mixed (Transformer Engine / MXFP4 path enabled).
• Optimizer: Distributed Optimizer (Adam) with overlap_param_gather=True.
• Memory: Selective Activation Checkpointing (~75% of layers).
• Learning Rate: Max LR 7e-6 (Phase 1) with Cosine Annealing and 3% Warmup.
• Data Loading: 32 workers, prefetch_factor=8, and persistent_workers=True.

Any idea what would be the best to do to increase GBS to accelerate the training process? I think that is the bottleneck, but not sure, maybe increasing TP?

[xXMrNidaXx]

Great setup! A few optimization paths to explore:

Why higher GBS slows down:
With EP=2, increasing GBS creates more communication overhead for expert synchronization. The MoE routing becomes a bottleneck.

Suggested changes:

1. Increase TP to reduce memory pressure

   tensor_model_parallel_size: 2  # Was 1
   pipeline_model_parallel_size: 1
   expert_model_parallel_size: 2
   # This allows higher MBS

2. Enable sequence parallelism

   sequence_parallel: true
   # Reduces activation memory, allows larger batches

3. Gradient accumulation instead of large GBS

   global_batch_size: 256
   micro_batch_size: 2
   # Accumulates 128 steps per GPU update
   # Same effective batch, less memory

4. FP8 for B200 (native support)

   fp8: true
   fp8_e4m3: true  # For forward
   fp8_hybrid: true  # Mixed precision

   B200 has excellent FP8 throughput.

5. Overlap more aggressively

   overlap_grad_reduce: true
   overlap_param_gather: true
   async_grad_allreduce: true

Quick benchmark:
Try TP=2, MBS=2, GBS=256 with gradient accumulation. Should unlock higher throughput without OOM.

We run large-scale NeMo training at Revolution AI — B200 FP8 path is key for 20B+ models. Happy to help debug further!