Skip to content

fix: OOM in deepscaler1.5b with sequence length = 16/24k #875

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
terrykong merged 5 commits into from
Aug 9, 2025
Merged

fix: OOM in deepscaler1.5b with sequence length = 16/24k #875

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
eee023b
update
soodoshll Aug 8, 2025
9040aa2
fix
soodoshll Aug 8, 2025
325b7b4
fix
soodoshll Aug 8, 2025
88c38b8
merge main
soodoshll Aug 8, 2025
988d2e9
update nemo
soodoshll Aug 8, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 0 additions & 4 deletions examples/configs/recipes/llm/grpo-deepscaler-1.5b-24K.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,3 @@ policy:
gpu_memory_utilization: 0.8
enforce_eager: True
max_model_len: ${policy.max_total_sequence_length}

cluster:
gpus_per_node: 8
num_nodes: 4
16 changes: 10 additions & 6 deletions nemo_rl/models/policy/dtensor_policy_worker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -720,6 +720,7 @@ def train(
logits = self.model.lm_head(outputs.last_hidden_state)
else:
logits = outputs.logits
del outputs

# Apply temperature scaling
logits = self._apply_temperature_scaling(logits)
Expand Down Expand Up @@ -786,6 +787,7 @@ def train(
global_valid_seqs,
global_valid_toks,
)
del logits

# skip the update for dummy batches
if mb_idx < iterator_len:
Expand Down Expand Up @@ -1044,17 +1046,19 @@ def get_logprobs(
placements=[Shard(sequence_dim), Shard(-1)],
)

logits = logits.to(torch.float32)
token_logprobs = get_logprobs_from_vocab_parallel_logits(
logits.to(torch.float32),
logits,
input_ids_dtensor,
seq_index_tensor,
)

assert token_logprobs.shape[1] == seq_len - 1
else:
if isinstance(logits, DTensor):
logits = logits.to(torch.float32)
token_logprobs = get_logprobs_from_vocab_parallel_logits(
logits.to(torch.float32), input_ids
logits, input_ids
)
else:
# Extract logprobs for each token in the sequence by gathering the logprob
Expand All @@ -1064,16 +1068,16 @@ def get_logprobs(
# token_ids: [batch_size, sequence_length] - actual tokens
# Output shape: [batch_size, sequence_length] - logprob of each token given previous
# We get logprob of token[t+1] from logits[t], prepending 0 to maintain sequence length

log_probs = torch.nn.functional.log_softmax(
outputs.logits.to(torch.float32), dim=-1
)
logits = outputs.logits.to(torch.float32)
log_probs = torch.nn.functional.log_softmax(logits, dim=-1)
next_tokens = input_ids[:, 1:]
log_probs = log_probs[:, :-1]
token_logprobs = log_probs.gather(
dim=-1, index=next_tokens.unsqueeze(-1)
).squeeze(-1)

del outputs, logits

token_logprobs = torch.cat(
[torch.zeros_like(token_logprobs[:, :1]), token_logprobs], dim=1
)
Expand Down