Skip to content

fix: align sequence_length_pad_multiple in lm_policy - #2182

Merged
terrykong merged 1 commit into
mainfrom
yukih/fix-make_sequence_length_divisible_by
Apr 1, 2026
Merged

fix: align sequence_length_pad_multiple in lm_policy#2182
terrykong merged 1 commit into
mainfrom
yukih/fix-make_sequence_length_divisible_by

Conversation

@yuki-97

@yuki-97 yuki-97 commented Mar 31, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Summary

Fix a crash in Megatron sequence packing when sequence_length_pad_multiple != make_sequence_length_divisible_by.

Problem

sequence_length_pad_multiple was computed as tp_size when cp_size=1, but pack_sequences actually pads each sequence to make_sequence_length_divisible_by = tp_size*2 which is defined in the config.
This mismatch causes the actual packed total to exceed pad_packed_seq_to, resulting in a crash at unpack time:

RuntimeError: The expanded size of the tensor (517) must match the existing size (502)

Fix

Replace cp_size * 2 * tp_size if cp_size > 1 else tp_size with config["make_sequence_length_divisible_by"] so both sides use the same padding factor. Safe for DTensor: when cp=1 bins become slightly more conservative but DTensor does not apply per-sequence alignment padding during packing.

Test plan

  • distillation-qwen3-32b-to-1.7b-base-1n4g-megatron-tp1pp2cp2-pack — previously crashed, now completes successfully.

@yuki-97
align sequence_length_pad_multiple in lm_policy
9dc3832 
Signed-off-by: Yuki Huang <yukih@nvidia.com>
@yuki-97
yuki-97 requested review from a team as code owners March 31, 2026 15:01
@copy-pr-bot

copy-pr-bot Bot commented Mar 31, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@yuki-97
yuki-97 requested a review from terrykong March 31, 2026 15:01
@yuki-97 yuki-97 added the CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) label Mar 31, 2026
@yuki-97

yuki-97 commented Mar 31, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 9dc3832

@terrykong
terrykong merged commit 56d8a0f into main Apr 1, 2026
29 of 30 checks passed
@terrykong
terrykong deleted the yukih/fix-make_sequence_length_divisible_by branch April 1, 2026 01:57
@pthombre pthombre mentioned this pull request Jul 31, 2026
Closed
pthombre added a commit that referenced this pull request Jul 31, 2026
@pthombre @claude
feat(sc): gate-authoritative token capture on the Gym tokidcap stack
1b3c83e 
Token-in/token-out capture for NeMo-Gym rollouts in the async
SingleController GRPO path, rebuilt on the nano SWE recipe branch
(3fcc696, main + the TQ-honoured SWE recipe) per
docs/design-docs/tq-gym-tokidcap-migration.md §9b.1a. The Gym submodule is
pinned to pthombre/tq-tokidcap-capture (b6051536 = upstream tokidcap stack
top 81ac2736/#2182 + the 7-commit gate/staging series; companion Gym PR
NVIDIA-NeMo/Gym#2278).

This squashes the capture series (S1-S5 primitives/worker/gate/receipts/
verification + the stack migration) into one commit ported onto the
main-lineage SC: the granular history lives on the prior branch head
624bb27. Highlights:

- TQTokenSink/TQTokenSource: per-call token deltas staged from the vLLM
  worker straight into TransferQueue, durable before the response releases;
  BlackboxFinalizer rebuilds digest-verified always-N training rows via
  Gym's terminal-aware linearize over run_builder.
- Identity: rollout ids ride the run body as the opaque _ng_rollout_id key
  (agents stamp /ng-rollout/<id>; all agent impls unmodified); the
  middleware-minted model_call_id is the call id; TQ sample ids ARE the
  capture keys.
- Gate hosting config: LineageIndex capacity derived from the training
  config; per-run control-plane bearer token; hard per-call control
  deadlines (gate death surfaces as failed dispatches + placeholders, not a
  silent retry stall); base capture dir under the run's log dir.
- Receipt-mode dispatch in RolloutManager (_generate_and_finalize:
  reserve-with-ids, seal -> receipt, finalize_group, commit_finalized,
  always-N with placeholders), weight-version fan-out to workers, gate
  metrics logging (token_in_rate, fallback-by-cause).
- Port adaptations onto main's structures: streaming run_rollouts (receipt
  postprocess branch + picklable aiohttp error re-raise), commit()'s
  pre-write and evicted-during-write guards keep main's semantics, legacy
  failure path keeps main's remove_group.

Dormant by default: everything is behind token_capture.enabled=false.
Working log: docs/design-docs/tq-gym-tokidcap-migration-log.md.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Pranav Prashant Thombre <pthombre@nvidia.com>
pthombre added a commit that referenced this pull request Aug 12, 2026
@pthombre @claude Pranav Thombre
feat(sc): gate-authoritative token capture on the Gym tokidcap stack
4c447d8 
Token-in/token-out capture for NeMo-Gym rollouts in the async
SingleController GRPO path, rebuilt on the nano SWE recipe branch
(3fcc696, main + the TQ-honoured SWE recipe) per
docs/design-docs/tq-gym-tokidcap-migration.md §9b.1a. The Gym submodule is
pinned to pthombre/tq-tokidcap-capture (b6051536 = upstream tokidcap stack
top 81ac2736/#2182 + the 7-commit gate/staging series; companion Gym PR
NVIDIA-NeMo/Gym#2278).

This squashes the capture series (S1-S5 primitives/worker/gate/receipts/
verification + the stack migration) into one commit ported onto the
main-lineage SC: the granular history lives on the prior branch head
624bb27. Highlights:

- TQTokenSink/TQTokenSource: per-call token deltas staged from the vLLM
  worker straight into TransferQueue, durable before the response releases;
  BlackboxFinalizer rebuilds digest-verified always-N training rows via
  Gym's terminal-aware linearize over run_builder.
- Identity: rollout ids ride the run body as the opaque _ng_rollout_id key
  (agents stamp /ng-rollout/<id>; all agent impls unmodified); the
  middleware-minted model_call_id is the call id; TQ sample ids ARE the
  capture keys.
- Gate hosting config: LineageIndex capacity derived from the training
  config; per-run control-plane bearer token; hard per-call control
  deadlines (gate death surfaces as failed dispatches + placeholders, not a
  silent retry stall); base capture dir under the run's log dir.
- Receipt-mode dispatch in RolloutManager (_generate_and_finalize:
  reserve-with-ids, seal -> receipt, finalize_group, commit_finalized,
  always-N with placeholders), weight-version fan-out to workers, gate
  metrics logging (token_in_rate, fallback-by-cause).
- Port adaptations onto main's structures: streaming run_rollouts (receipt
  postprocess branch + picklable aiohttp error re-raise), commit()'s
  pre-write and evicted-during-write guards keep main's semantics, legacy
  failure path keeps main's remove_group.

Dormant by default: everything is behind token_capture.enabled=false.
Working log: docs/design-docs/tq-gym-tokidcap-migration-log.md.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Pranav Prashant Thombre <pthombre@nvidia.com>
pthombre added a commit that referenced this pull request Aug 12, 2026
@pthombre @claude Pranav Thombre
feat(sc): gate-authoritative token capture on the Gym tokidcap stack
9b5bdf6 
Token-in/token-out capture for NeMo-Gym rollouts in the async
SingleController GRPO path, rebuilt on the nano SWE recipe branch
(3fcc696, main + the TQ-honoured SWE recipe) per
docs/design-docs/tq-gym-tokidcap-migration.md §9b.1a. The Gym submodule is
pinned to pthombre/tq-tokidcap-capture (b6051536 = upstream tokidcap stack
top 81ac2736/#2182 + the 7-commit gate/staging series; companion Gym PR
NVIDIA-NeMo/Gym#2278).

This squashes the capture series (S1-S5 primitives/worker/gate/receipts/
verification + the stack migration) into one commit ported onto the
main-lineage SC: the granular history lives on the prior branch head
624bb27. Highlights:

- TQTokenSink/TQTokenSource: per-call token deltas staged from the vLLM
  worker straight into TransferQueue, durable before the response releases;
  BlackboxFinalizer rebuilds digest-verified always-N training rows via
  Gym's terminal-aware linearize over run_builder.
- Identity: rollout ids ride the run body as the opaque _ng_rollout_id key
  (agents stamp /ng-rollout/<id>; all agent impls unmodified); the
  middleware-minted model_call_id is the call id; TQ sample ids ARE the
  capture keys.
- Gate hosting config: LineageIndex capacity derived from the training
  config; per-run control-plane bearer token; hard per-call control
  deadlines (gate death surfaces as failed dispatches + placeholders, not a
  silent retry stall); base capture dir under the run's log dir.
- Receipt-mode dispatch in RolloutManager (_generate_and_finalize:
  reserve-with-ids, seal -> receipt, finalize_group, commit_finalized,
  always-N with placeholders), weight-version fan-out to workers, gate
  metrics logging (token_in_rate, fallback-by-cause).
- Port adaptations onto main's structures: streaming run_rollouts (receipt
  postprocess branch + picklable aiohttp error re-raise), commit()'s
  pre-write and evicted-during-write guards keep main's semantics, legacy
  failure path keeps main's remove_group.

Dormant by default: everything is behind token_capture.enabled=false.
Working log: docs/design-docs/tq-gym-tokidcap-migration-log.md.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Pranav Prashant Thombre <pthombre@nvidia.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@terrykong terrykong terrykong approved these changes

Assignees

No one assigned

Labels

CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version)

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@yuki-97 @terrykong