THUDM / THUDM/slime

OOM when training with long sequences despite using dynamic batch size and sequence parallel

Open
#1,522 4 comments 0 reactions 0 assignees View on GitHub

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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.