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?