OOM when training with long sequences despite using dynamic batch size and sequence parallel
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 8.5k
- Forks
- 1.3k
- Avg merge
- 5h 36m
- Merged PRs (30d)
- 22
Description
Problem
I need to configure Megatron for training a 4B model with 64K max-response-len using PPO on 4x H200 GPUs.
I'm experiencing Out of Memory (OOM) errors when training with 64K sequences, even with --use-dynamic-batch-size and --sequence-parallel enabled.
My Understanding
I set --max-tokens-per-gpu . I expected:
- Dynamic batch size would automatically adjust micro_batch_size based on sequence length
- This would prevent activations from exceeding GPU memory
- optimizer.step() memory usage should remain stable regardless of global batch size
However: It seems unlikely to avoid OOM by simply reducing global batch size, since optimizer memory is independent of global batch size.
MyConfig
TP_SIZE=2
PP_SIZE=1
CP_SIZE=1
EP_SIZE=1
ETP_SIZE=1
MAX_LEN=$((1024 * 64)) # 64K
MAX_TOKENS_PER_GPU=$((($MAX_LEN / $CP_SIZE) + 1024)) # ~65K
ROLLOUT_BATCH_SIZE=16
N_SAMPLES_PER_PROMPT=4
NUM_STEPS_PER_ROLLOUT=4
--rollout-batch-size $ROLLOUT_BATCH_SIZE \
--n-samples-per-prompt $N_SAMPLES_PER_PROMPT \
--rollout-max-response-len $MAX_LEN \
--colocate \
--actor-num-gpus-per-node 2 \
--tensor-model-parallel-size 2 \
--sequence-parallel \
--pipeline-model-parallel-size 1 \
--context-parallel-size 1 \
--recompute-granularity full \
--recompute-method uniform \
--recompute-num-layers 1 \
--use-dynamic-batch-size \
--max-tokens-per-gpu 65536 \
--transformer-impl transformer_engine \
--bf16 \
--fp8-format e4m3 \
--fp8-recipe blockwise
Is there any way to avoid the OOM problem during long sequence training?
Any guidance appreciated! 🙏
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
No source file or test is named. Start by reproducing the supplied 64K configuration, then trace the --use-dynamic-batch-size and --max-tokens-per-gpu paths through training and optimizer.step() while measuring memory. Done means identifying the OOM cause and documenting or validating a configuration or code change that handles long-sequence training.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100