Skip to content

[FIX] Remove hardcoded seq len #497

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

Jump to bottom
Merged
merged 2 commits into from
Oct 23, 2025
Merged

[FIX] Remove hardcoded seq len #497

merged 2 commits into from
Oct 23, 2025

Conversation

@felipemello1
Copy link
Contributor

@felipemello1 felipemello1 commented Oct 23, 2025 •
edited
Loading

fixes #495

image

Summary

Fixed rope_cache size mismatch in GRPO training.

Problem

The trainer.training.seq_len parameter was hardcoded to 2048, causing an AssertionError when the actual sequence length (max_req_tokens + max_res_tokens) exceeded this value.

Traceback (most recent call last):
  File "/fsx/lewis/miniconda3/envs/forge/lib/python3.12/site-packages/monarch/_src/actor/actor_mesh.py", line 932, in handle
    result = await the_method(*args, **kwargs)
  File "/fsx/lewis/git/torchforge/src/forge/actors/reference_model.py", line 181, in forward
    logits = self.model(input_ids)
  File "/fsx/lewis/miniconda3/envs/forge/lib/python3.12/site-packages/torchtitan/models/qwen3/model/model.py", line 491, in forward
    h = layer(h, self.rope_cache, attention_masks)
  File "/fsx/lewis/miniconda3/envs/forge/lib/python3.12/site-packages/torchtitan/models/qwen3/model/model.py", line 345, in forward
    x = x + self.attention(self.attention_norm(x), rope_cache, attention_masks)
  File "/fsx/lewis/miniconda3/envs/forge/lib/python3.12/site-packages/torchtitan/models/qwen3/model/model.py", line 219, in forward
    xq, xk = apply_rotary_emb(xq, xk, rope_cache)
  File "/fsx/lewis/miniconda3/envs/forge/lib/python3.12/site-packages/torchtitan/models/qwen3/model/model.py", line 91, in apply_rotary_emb
    rope_cache = reshape_for_broadcast(rope_cache, xq)
  File "/fsx/lewis/miniconda3/envs/forge/lib/python3.12/site-packages/torchtitan/models/qwen3/model/model.py", line 79, in reshape_for_broadcast
    assert rope_cache.shape == (seqlen, head_dim * 2)
AssertionError

Root Cause

qwen3_1_7b.yaml (seq_len=2048, max_req_tokens=512, max_res_tokens=16384)
  → Qwen3Model.__init__() (precomputes rope_cache with size 2048)
  → model.forward() (receives 16896 tokens)
  → reshape_for_broadcast() (assertion fails: 2048 ≠ 16896)

Solution

Just set seq_len to the sum of max_res_tokens + max_req_tokens

training:
    seq_len: ${sum:${max_req_tokens},${max_res_tokens}}

To support "sum" in the yaml, i had to register the op with OmegaConf

@meta-cla meta-cla bot added the CLA Signed This label is managed by the Meta Open Source bot. label Oct 23, 2025
@felipemello1 felipemello1 mentioned this pull request Oct 23, 2025
Closed
@felipemello1 felipemello1 changed the title [draft] [FIX] Enable larger seq len [FIX] Remove hardcoded seq len Oct 23, 2025
allenwang28
allenwang28 approved these changes Oct 23, 2025
View reviewed changes
Copy link
Contributor

@allenwang28 allenwang28 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thanks for the quick fix!

@joecummings joecummings merged commit 6919943 into meta-pytorch:main Oct 23, 2025
8 checks passed
@felipemello1 felipemello1 deleted the fix_seq_len branch October 23, 2025 16:46
photomz pushed a commit to photomz/forge that referenced this pull request Oct 25, 2025
@felipemello1
[FIX] Remove hardcoded seq len (meta-pytorch#497)
374a8fa 
Co-authored-by: Felipe Mello <[email protected]>
DNXie pushed a commit to DNXie/forge that referenced this pull request Nov 5, 2025
@felipemello1 @DNXie
[FIX] Remove hardcoded seq len (meta-pytorch#497)
d16f511 
Co-authored-by: Felipe Mello <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@allenwang28 allenwang28 allenwang28 approved these changes

@casteryh casteryh Awaiting requested review from casteryh

Assignees

No one assigned

Labels

CLA Signed This label is managed by the Meta Open Source bot.

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

[GRPO app] Example fails when max_res_tokens: 16384 with assert rope_cache.shape == (seqlen, head_dim * 2)

3 participants

@felipemello1 @allenwang28 @joecummings