[FIX] Remove hardcoded seq len

[felipemello1]

fixes #495

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

[allenwang28]

thanks for the quick fix!