Reduce CPU-side materialization overhead in the policy-only training pipeline

[CAICAIIs]

Motivation

The policy-only trainer (GRPO/GSPO) spends its per-step CPU time between rollout and train building and repacking Python lists. For each prompt group the current data flow materializes response data three times as Python lists before it reaches the CUDA training tensors:

1. rollout logprobs are converted GPU→CPU per sequence (engine/worker.py, then backend.py rollout.logprobs[index, : len(tokens)].tolist());
2. PolicyOnlyTrainer._materialize_train_batch decodes each completion individually, allocates per-sample [0.0] * prefix + resp logprob/loss-mask lists inside make_reward_record, and builds a TrainSequence with four freshly allocated per-token lists (prompt_mask, tokens, logprobs, advantages);
3. make_train_pack re-iterates those lists field by field (pad_rows, _make_prompt_mask, _make_advantages) into right-padded tensors.

With long-CoT RL (large max_new_tokens, batch_size × n_samples completions), this is O(response tokens) of Python list allocation and iteration executed serially while the GPUs sit idle between rollout and the next backend.train call. It is also the unit of work the experimental async trainer (#487) will queue, so reducing this overhead benefits both the synchronous and the asynchronous paths.

Measured baseline (single-node 8×A100, GRPO + GSM8K + Qwen3-0.6B, --attn-backend native)

Per-step phase timing from train_stats (step_e2e_time_s = step_rollout_time_s + step_train_time_s + materialization):

workload	completions/step	avg response tokens	rollout	train	materialization	% of step
small (max-new-tokens 512, batch 4 × n 4)	16	512	8.27 s	2.42 s	0.06 s	0.6 %
mid (max-new-tokens 1024, batch 8 × n 4)	32	842	37.54 s	7.33 s	0.33 s	0.7 %

This is roughly 12 µs of pure Python list work per response token. The per-step materialization cost is linear in response tokens, so it scales to the workloads RL post-training is increasingly dominated by: at ~1M response tokens per step (e.g. 8 GPUs, batch 32 × n 8 × max-new-tokens 4096), the same path would spend an estimated 12–13 s/step in serial CPU list building while all GPUs idle — a double-digit share of the step and growing with CoT length.

Also observed on the same runs: ratio_mean=1.0, ratio_std=0.0 for the current unit-valued GRPO surrogate, consistent with the discussion in #68.

Environment: 8×A100-40GB, torch 2.13.0+cu130, AReno main @ 9218129.

Proposed optimization

A behavior-preserving refactor of the rollout→train materialization path, scoped to the CUDA policy-only pipeline:

• replace per-completion tokenizer.decode calls with one batched decode per prompt batch where the tokenizer supports it;
• build the training pack directly from rollout tensors where they are still on-device, or from a single flattened CPU pass, instead of allocating per-token TrainSequence lists and repacking them a second time;
• compute group-relative advantages (compute_group_advantages) with one vectorized pass over the whole batch using prompt-group boundaries, instead of one numpy call per group;
• drop or amortize the per-step float(np.mean(rollout_logprobs)) over the full response-token set (keep the metric, avoid the full-token reduction when it only feeds a log line);
• keep the public TrainSequence model and all loss-function contracts unchanged; only internal layout/repack helpers may change.

The CLI surface is unchanged:

areno train --algo grpo --ckpt <model> --dataset-path gsm8k:main \
  --reward-fn-path examples/math/math_verify_reward.py \
  --batch-size 8 --n-samples 4 --max-new-tokens 1024

Validation plan

• CPU tests asserting the refactor is behavior-preserving: identical loss and gradient values for GRPO and GSPO on packed and padded batches before/after (reuse the existing tests/ -k cpu fixtures). Baseline already collected: tests/test_algorithms_cpu.py tests/test_losses_rewards_cpu.py tests/test_more_losses_cpu.py tests/test_agentic_cpu.py tests/test_config_data_cpu.py tests/test_parallel_partition_cpu.py → 135 passed;
• a bounded GPU end-to-end run on a long-CoT workload (max-new-tokens ≥ 1024) comparing per-step wall time, rollout→train gap, and GPU utilization before/after; report both numbers in the PR. The pre-change baseline above (small + mid workloads on 8×A100) is reproducible from areno train --algo grpo --ckpt Qwen/Qwen3-0.6B --dataset-path gsm8k:main --dataset-loader-fn examples/math/dataset_loader.py --reward-fn-path examples/math/math_verify_reward.py --attn-backend native;
• a short profile excerpt (e.g. py-spy / cProfile) showing where the CPU time moves after the change.

I have access to a single-node 8×A100 host and can run the bounded GPU end-to-end validation for this change.

Questions for maintainers

1. Is a TrainSequence-level layout change (e.g. carrying flat tensors instead of per-token lists) acceptable inside the backend, as long as the public model and loss contracts stay unchanged?
2. Should the group-advantage vectorization live in areno/api/rewards.py (shared with PPO/DPO paths) or stay trainer-local?
3. Would maintainers prefer this as one cohesive refactor PR, or split the logging amortization into a separate small PR?

Alternatives considered

• Eliminating the GPU→CPU transfer entirely would require scoring rewards on-device, which is impossible for arbitrary Python reward_fns and is the separate concern of feature: add experimental async policy trainer #487 (async overlap).
• Optimizing only the reward hooks (per Identify slow reward hooks and samples #242) leaves the materialization and repack overhead untouched.
• Leaving the path as-is keeps the code simple but forfeits measurable per-step latency on long-CoT workloads, which is the workload RL post-training is increasingly dominated by.

[CAICAIIs]

Hi, I'd like to claim this issue , please assign it to me.

[xsuler]

Hi @CAICAIIs, this is a very good proposal. Just note that the main loop of the algorithm should remain CPU-only in your implementation, and you may implement amortized statistics over rollout logprobs instead of dropping them. The group-advantage vectorization should live in a separate file like advantages.py rather than being placed in rewards.py. A single cohesive refactor PR will be preferred, along with a complete performance enhancement report.

[CAICAIIs]

Performance validation on merged main (#508 + #516)

Consolidated measurement on the merged main (8×A100-40GB, torch 2.13.0+cu130,
GRPO + GSM8K + Qwen3-0.6B, --attn-backend native, single GPU).

End-to-end materialization (step_e2e − step_rollout − step_train)

workload	tokens/step	before (#506 proposal)	after (merged main)
mid (b8×n4, max-new-tokens 1024)	~25K	0.327 s	0.17–0.19 s (2 steps)
long-CoT (b2×n4, max-new-tokens 4096)	~6–9K	—	0.06–0.13 s (2 steps, peak mem 21.7 GB)

The apples-to-apples mid workload shows materialization down ~45% (0.327 → ~0.18 s);
rollout/train phases fluctuate with shared-host load, so the pack-path
micro-benchmarks below are the controlled measurements.

Cumulative micro-benchmarks (identical random data, same host)

path	original	after #508+#516	speedup
make_train_pack (52K tokens)	14.27 ms	5.10 ms	2.8×
rollout-logprob statistics (n=1M)	46.99 ms	4.11 ms	11.4×, O(1) memory
group advantages (G=128 × n=16)	3.65 ms	0.16 ms	23.5×

Safety note

The long-CoT run (b2×n4×4096) peaked at 21.7 GB / 40 GB; a b8×n4×4096 attempt
was stopped by a 36 GB guard before any actual OOM. All runs on this shared
host were process-isolated.

Status

• perf(engine): reduce CPU-side materialization overhead in the policy-only trainer #508 (materialize refactor): merged
• perf(engine): vectorize train-pack construction #516 (train-pack vectorization): merged
• fix(engine): allocate distributed init ports outside the ephemeral range #518 (coordinator-held rendezvous store, fixes a startup EADDRINUSE found
  during this validation): open for review

Remaining in-scope costs: reward-fn scoring (per #242) and the remaining
pack path after the Q1 layout decision.