[rl] Remove RL H100 CI to A10G and drop oversized sample before packing by wwwjn · Pull Request #3855 · pytorch/torchtitan

wwwjn · 2026-07-01T23:24:31Z

RL H100 failing: https://github.com/pytorch/torchtitan/actions/runs/28451614145/job/84315141822#step:10:1576

Also move the H100 test to A10G to solve the failing permission issue in H100

tianyu-l · 2026-07-07T05:21:14Z

    def collate(rows: list[dict]) -> TrainingMicrobatch:
        """Concatenate packed rows into a single ``[B, L]`` TrainingMicrobatch."""
        return TrainingMicrobatch(
            token_ids=torch.cat([row["input_ids"] for row in rows]),


According to https://docs.pytorch.org/docs/2.12/generated/torch.cat.html:

All tensors must either have the same shape (except in the concatenating dimension) or be a 1-D empty tensor with size (0,).

It sounds that without this PR it should just error out?

Yes you are right, this is the exact failure that we saw in H100 https://github.com/pytorch/torchtitan/actions/runs/28451614145/job/84315141822#step:10:1576 which haven't been fixed yet. After this PR fixed the torch.cat() issue, we also see another file write permission issue, which is separate

tianyu-l · 2026-07-07T05:25:27Z

+            logger.warning(
+                "Batcher dropped %d/%d training sample(s) exceeding seq_len=%d; "
+                "increase seq_len or reduce generation length to keep them.",
+                num_dropped,
+                len(training_sample_group.training_samples),
+                self.seq_len,
+            )


I'm worried that this will be verbose.

btw I think we should improve this greedy batcher, see

TODO(async-rl): assignment is greedy next-fit. Swap in smarter algorithms here -- e.g. best-fit

I can take the task of updating the packing algorithm. Will create a separate PR to update to n^2 based packing

But no matter which algorithm, we still need to truncate or drop and handle the corner case that one example is too long for trainining

tianyu-l · 2026-07-07T05:27:35Z

            if current_row and current_len + num_tokens_to_pack > self.seq_len:
                rows.append(current_row)
                current_row, current_len = [], 0

            current_row.append(training_sample)


I think we should move the filter logic to here, instead of doing a "preprocessing" step.

The training_samples are returned from _take_groups_for_train_step, which will yield num_groups_per_train_step samples and pack into rows here. If we drop here, trainer will receive less than num_groups_per_train_step samples. The worse case we will drop all and trainer recieve 0 rows.

The trigger that decides when to pack lives in add_training_samples(), and if we don't want to count the super long samples into num_groups_per_train_step, we need to filter it out earlier than later

Fold the H100 RL integration workflow into the A10G (g5.48xlarge) workflow and delete the H100-specific one. The 0.6B model at TP=4 for both trainer and generator OOMs on A10G under batch-invariant mode (the batch-invariant kernels need extra memory), so the batch-invariant E2E test and the loss guard now run the tiny qwen3 debugmodel at TP=4. drop_zero_std_reward_groups=False keeps the zero-reward groups so the trainer forward still runs on random-init completions. Bitwise-parity tests stay on 0.6B at TP=2 (forward-only, no OOM). Also fix loss_compare.py to pass --async-loop.* (the config field is async_loop; --training-loop.* was silently wrong). Note: the deterministic loss golden (rl_grpo_cuda.txt) must be regenerated on A10G for the debug config before the loss-guard step passes.

The varlen batch-invariant config is tp=4 but the CI launched the test at --nproc-per-node=2, so ParallelDims asserted (tp(4) != WORLD_SIZE(2)). Launch varlen with 4 procs; flex stays at 2 (its config is tp=2).

…10G budget) 0.6B TP=4 does not OOM on A10G, but the full 0.6B TP=4 E2E + loss-guard generation overruns the 90-min job timeout, and the loss-guard golden (rl_grpo_cuda.txt) holds the debug config's near-zero losses so the 0.6B run would also fail --assert-equal. Revert these two heavy vehicles to the tiny debug batch-invariant config (fast, golden-matching). The bitwise parity tests keep the real 0.6B model (varlen TP=4, flex TP=2) for meaningful parity coverage.

The loss guard used --dump-folder $RUNNER_TEMP/rl_loss_guard, but only $RUNNER_TEMP/artifacts-to-be-uploaded is chowned writable for the container user (setup step), so train.py hit PermissionError creating the dir. Point it inside artifacts-to-be-uploaded (also uploads the traces).

The golden (rl_grpo_cuda.txt) held small non-zero losses from before #3802, which set drop_zero_std_reward_groups=False for debug configs. With that + a random-init model, every reward group is zero-std -> zero advantage -> the GRPO loss is deterministically 0.0 (all 10 steps, confirmed in CI). Update the golden to 0.0 (arch-independent) and document why in the workflow comment.

The loss guard used a random-init model (--hf-assets-path tests/assets/tokenizer), so all rewards were 0 -> zero-std groups -> zero advantage -> loss==0.0. Point it at the real 0.6B ($MODEL_PATH) so completions earn varied rewards and the loss guard has meaningful non-zero values again. Lower job timeout to 60 (debug suite runs ~40 min). Golden will be regenerated on A10G from this run's actual losses.

The loaded-0.6B loss guard crashed at the generator's distributed init with EADDRINUSE (port already in use) because it ran after the E2E suite, which left a port bound. Run it first, in a clean container, so its generator gets a free port. (This capture run will fail the assert against the placeholder golden and print the real A10G 0.6B losses, which become the new golden.)

Captured from the loss-guard-first CI run (run 28961347634): the loaded-weights 0.6B loss guard produces meaningful non-zero, full-precision losses (varied rewards -> non-zero advantage). Deterministic + batch-invariant + seed=42 makes these reproducible across A10G runs.

…steps=0) The loss guard was reproducible for steps 1-9 but the last step diverged run-to- run: with max_offpolicy_steps=3 the tail step's rollouts depend on async generator/trainer timing (variable on A10G). Set it to 0 (lockstep, fully on-policy) so the loss curve is deterministic -- also matches the config's 'true on-policy' docstring. Golden will be re-captured for the on-policy curve.

Move max_offpolicy_steps=0 out of the shared 0.6B config and into loss_compare.py's train invocation (next to its other determinism overrides), so only the loss guard runs lockstep/on-policy for a deterministic curve; the config keeps its off-policy production default. Golden re-captured for the on-policy curve (run 28964105953).

…s visible The trainer-vs-generator logprob diff (0 under batch-invariant mode) is computed every step by DAPOLoss/GRPOLoss but was never surfaced: it only went to TensorBoard/W&B backends (both off by default), and was not in console_log_keys_train. So the E2E batch-invariant test computed the signal but dropped it. Add it to the console keys so zero logprob diff is visible in the logs (incl. CI) with no backend enabled.

…nfig Point the E2E batch-invariant test at rl_grpo_qwen3_0_6b_varlen_batch_invariant with loaded weights (suite --hf_assets_path) and on-policy (--async-loop.max-offpolicy-steps 0). The debug config gave a weak check: random init + drop_zero_std=False -> zero gradient -> weights never change, so the diff was trivially 0. Real 0.6B updates weights each step, so bit_wise/logprob_diff/max == 0 now verifies kernel batch-invariance AND weight sync. Removes the now-unused rl_grpo_qwen3_debug_varlen_batch_invariant config.

torchstore logs '[ts-transport] resolved=...' at INFO on every weight-sync transport resolve, per generator/trainer rank -- floods the CI logs. Raise the torchstore.transport logger to WARNING in both RL actors (after init_logger).

tianyu-l · 2026-07-09T03:47:04Z

+        # Quiet torchstore's per-op transport-resolve INFO spam (very noisy in CI).
+        logging.getLogger("torchstore.transport").setLevel(logging.WARNING)


cc @meetv18

This is because I saw CI prints has a lot of torchstore related INFO print, which makes the whole logging hard to read. @meetv18 is it ok what we set the level to WARNING? Do you need these info for debugging?

ahh yes, we should change this.

tianyu-l · 2026-07-09T03:48:14Z

@@ -859,7 +859,7 @@ def rl_grpo_qwen3_0_6b_varlen_batch_invariant() -> Controller.Config:
            training=TrainingConfig(),


nit: can remove this line

tianyu-l · 2026-07-09T03:48:44Z

            parallelism=ParallelismConfig(
                data_parallel_shard_degree=1,
-                tensor_parallel_degree=2,
+                tensor_parallel_degree=4,


should default to trainer FSDP 2, 3 generators each TP 2

0.6B model is used in CI, which A10G only have ~20G GPU memory. If we set TP=2 for 0.6B model + batch_invariant mode ON, it will OOMed on A10G. I will change the Batch-invariant config to trainer TP2, 3 generator TP2, and override in integration_tests.py

TrainingConfig() is already the default (field default_factory=TrainingConfig), so the explicit arg is a no-op. Remove it (and its inline comment, which duplicated the config docstring's fp32-master/FSDP-bf16 parity explanation) in rl_grpo_qwen3_0_6b_varlen_batch_invariant.

Switch the trainer from TP=2 to FSDP=2 (dp_shard=2, tp=1); generators stay at 3x TP=2 (8 GPUs total). More standard trainer parallelism for the non-BI functional test. Rename tp2_no_compile -> fsdp2_gen_tp2_no_compile.

wwwjn · 2026-07-09T16:54:35Z

+        # Quiet torchstore's per-op transport-resolve INFO spam (very noisy in CI).
+        logging.getLogger("torchstore.transport").setLevel(logging.WARNING)


This is because I saw CI prints has a lot of torchstore related INFO print, which makes the whole logging hard to read. @meetv18 is it ok what we set the level to WARNING? Do you need these info for debugging?

wwwjn · 2026-07-09T16:59:21Z

            parallelism=ParallelismConfig(
                data_parallel_shard_degree=1,
-                tensor_parallel_degree=2,
+                tensor_parallel_degree=4,


0.6B model is used in CI, which A10G only have ~20G GPU memory. If we set TP=2 for 0.6B model + batch_invariant mode ON, it will OOMed on A10G. I will change the Batch-invariant config to trainer TP2, 3 generator TP2, and override in integration_tests.py

…I overrides to TP=4 Make rl_grpo_qwen3_0_6b_varlen_batch_invariant document the canonical batch-invariant setup: trainer TP=2 + 3 generators TP=2 (matched TP keeps bitwise trainer/generator parity; 3 engines for generation throughput). A10G CI (~20GB/GPU) OOMs with TP=2 + batch-invariant in the heavy full-loop runs, so the memory-heavy runs override to trainer TP=4 + 1 generator TP=4 (shards more per GPU; matches the arch-specific golden generated at TP=4): - E2E integration test overrides parallelism in-command. - loss_compare.py now forwards unknown CLI args to train.py, so the loss guard pins TP=4 without touching the shared config or the golden. - bitwise parity varlen test runs at the config's TP=2 (light single-forward comparison; nproc 4->2) to verify the canonical setup's batch invariance. Also revert the non-BI E2E test back to trainer TP=2 (from the FSDP=2 attempt).

…mpare.py convention) Replace the parse_known_args passthrough with an explicit --options string argument, mirroring --baseline-options/--test-options in scripts/loss_compare.py. The yaml pins the loss guard to TP=4 via --options='...'.

The RL loss_compare builds an argv list for subprocess (no shell), so the --options string is tokenized in Python. Our overrides are plain --flag value pairs, so str.split() suffices; matches the core scripts/loss_compare.py which has no shlex dependency.

…iant pin The FA2 auto-num_splits NaN with paged KV (pytorch/pytorch#179760) was fixed upstream (Dao-AILab/flash-attention#2542, pulled into PyTorch via pytorch/pytorch#185281), so the num_splits=1 force for FA2 is no longer needed and pessimized split-k parallelism. Keep num_splits=1 only in batch-invariant mode (off ROCm), which still requires it for deterministic split-k reductions. Applies to both the core varlen path and the RL vLLM paged-KV decode path.

is_in_batch_invariant_mode() is only ever set by the RL path, and the RL controller (__post_init__) already errors on batch_invariant + ROCm before any forward runs, so the torch.version.hip guard in the attention forward is unreachable. Remove it (and the now-obvious comment) from both the core varlen path and the RL vLLM decode path.

The #3235 fix is NOT in the PyTorch nightly: pytorch/pytorch#185281 (the flash submodule bump meant to carry Dao-AILab/flash-attention#2542) was CLOSED without merging on 2026-05-27, and no other bump landed it. #3235 was closed prematurely. CI proof: with the workaround removed, the non-BI E2E test hit a CUDA device-side assert in generator_2's FA2 paged-KV decode (self._engine.step -> ncclAllReduce -> unhandled cuda error) at step 5 after steps 1-4 decoded fine -- the exact intermittent corruption profile of #3235. No EADDRINUSE this run (not the port race); the earlier green re-run with num_splits=1 passed. Restore the workaround until the upstream fix actually lands in the nightly.

pytorch/pytorch#185281 DID land via pytorchmergebot (commit ae9ca07); the CLOSED/merged=null state is a bot rebase-merge artifact, not a failed merge. So the flash-attention fix is in the nightly and num_splits=auto is safe. Reverts the earlier over-cautious restore (ff219d8).

…tor port race) Concurrently-spawned generators on this_host() each auto-pick a torch.distributed rendezvous port; two racing to the same one surface as EADDRINUSE at init, or (when a late binder joins a peer's store) cross-wired process groups that corrupt collectives and trip a device-side assert mid-decode. Assign each mesh a distinct free MASTER_PORT via the launch env.

Removes the port helper from the launch path (train.py back to main). This also drops train.py from the PR's changed files, clearing the ufmt lint that only fired because touching train.py surfaced a pre-existing _preimport_torch docstring nit. The multi-generator single-host port flake will be handled in the CI test harness instead.

…race

test 1 (no compile) and test 2 (compile) now run trainer FSDP=2 (dp_shard=2, tp=1) + 3 generators TP=2. Restore both attention.py to main (num_splits=1 workaround intact) so this PR is FA2-neutral (#3235 cleanup is a separate PR). This run checks whether the step-5 OOB gather recurs with the workaround in place -- expected gone if it was num_splits=auto x multi-generator.

Merged 4b041be ([rl] Remove RL H100 CI to A10G + drop oversized sample before packing, pytorch#3855) from upstream/main. Entirely experiments/rl/ + CI YAML; no models/{llama3,deepseek_v3,qwen3} changes, agpt/moe diff empty, zero conflicts. No replay, no smoke required (nothing our runs touch).

…e RL H100 badge - Comment out TestBitwiseParityFlex torchrun step: the flex generator hits the vLLM flex_attention_compiled recompile-limit deadlock (#3898), which hangs the A10G RL integration job until CI times out. Disabled to confirm flex is the hang source (varlen teardown was fixed separately). - Remove the dead RL H100 CI badge from experiments/README.md (the workflow was removed in #3855); it 404s and fails the lychee link checker.

pytorch-bot Bot added ciflow/8gpu ciflow/rl labels Jul 1, 2026

meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Jul 1, 2026

wwwjn requested review from fegin, tianyu-l and wconstab as code owners July 2, 2026 16:27

wwwjn force-pushed the fix-rl-h100 branch from d3b2fc7 to a2518f2 Compare July 2, 2026 16:44

tianyu-l requested changes Jul 7, 2026

View reviewed changes

wwwjn changed the title ~~[rl] Fix RL H100 CI: drop oversized sample before packing~~ [rl] Remove RL H100 CI to A10G and drop oversized sample before packing Jul 7, 2026

wwwjn changed the title ~~[rl] Remove RL H100 CI to A10G and drop oversized sample before packing~~ [WIP][rl] Remove RL H100 CI to A10G and drop oversized sample before packing Jul 7, 2026

wwwjn added ciflow/rl and removed ciflow/rl labels Jul 8, 2026

wwwjn added 16 commits July 8, 2026 17:16

update

76b1c2d

update

c1f266d

update

2f26036

update

e229c23

[rl] bitwise parity CI: launch varlen at nproc=4 to match tp=4 config

fcfe7ea

The varlen batch-invariant config is tp=4 but the CI launched the test at --nproc-per-node=2, so ParallelDims asserted (tp(4) != WORLD_SIZE(2)). Launch varlen with 4 procs; flex stays at 2 (its config is tp=2).

[rl] simplify E2E batch-invariant test comment

02eda28

wwwjn force-pushed the fix-rl-h100 branch from e40d3e4 to 02eda28 Compare July 9, 2026 00:16

wwwjn changed the title ~~[WIP][rl] Remove RL H100 CI to A10G and drop oversized sample before packing~~ [rl] Remove RL H100 CI to A10G and drop oversized sample before packing Jul 9, 2026

tianyu-l approved these changes Jul 9, 2026

View reviewed changes

wwwjn added 2 commits July 9, 2026 09:58

[rl] E2E non-BI 0.6B test: trainer FSDP=2 (was TP=2), 3 gens TP=2

9b03374

Switch the trainer from TP=2 to FSDP=2 (dp_shard=2, tp=1); generators stay at 3x TP=2 (8 GPUs total). More standard trainer parallelism for the non-BI functional test. Rename tp2_no_compile -> fsdp2_gen_tp2_no_compile.

wwwjn commented Jul 9, 2026

View reviewed changes

wwwjn added 6 commits July 9, 2026 10:14

[rl] simplify loss-guard TP=4-override comment

39ba254

[rl] shorten num_splits batch-invariant comment

7f44ea9

tianyu-l reviewed Jul 9, 2026

View reviewed changes

Comment thread torchtitan/experiments/rl/models/attention.py Outdated

Comment thread torchtitan/models/common/attention.py Outdated

[rl] restore ROCm rationale in num_splits comment

6014daa

wwwjn mentioned this pull request Jul 9, 2026

[rl] FAv2 + num_splits=auto + paged kv produced NaN during inference #3235

Open

wwwjn added 9 commits July 9, 2026 12:13

[rl] shorten loss_compare sys.path import comment

3ffde07

[rl] trim A10G comment in E2E batch-invariant test

514750f

[rl] PROBE: drop test 1 to 1 generator to isolate num_splits vs port …

bdf0925

…race

wwwjn merged commit 4b041be into main Jul 9, 2026
19 of 22 checks passed

		# Quiet torchstore's per-op transport-resolve INFO spam (very noisy in CI).
		logging.getLogger("torchstore.transport").setLevel(logging.WARNING)

		@@ -859,7 +859,7 @@ def rl_grpo_qwen3_0_6b_varlen_batch_invariant() -> Controller.Config:
		training=TrainingConfig(),

Uh oh!

Conversation

wwwjn commented Jul 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

wwwjn Jul 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

wwwjn commented Jul 1, 2026 •

edited

Loading

wwwjn Jul 7, 2026 •

edited

Loading