Fix Prime-RL custom trainer SFT staging

[bingran-you]

Summary

• replace the Mobile300 Prime-RL compatibility path's decoded token-suffix/passthrough-template staging with pre-tokenized custom-trainer tensors
• add a run-local Prime-RL data shim that feeds input_ids / target_ids / loss_mask / position_ids directly to SFTDataset._process
• keep the sample-mean loss shim, normalize the old custom-trainer-token-suffix spelling to custom-trainer-pretokenized, and update the CLI/help/tests

Why

The previous BenchFlow wrapper preserved the historical custom trainer's token tail, but then decoded it into a fake assistant message. With Prime-RL assistant-only labeling, that changed the old full-label custom trainer objective and could train on the wrong role/message semantics. The compatibility profile now stages the same shifted tensor targets the historical trainer exposed to HF CausalLM loss, while leaving Prime-RL optimizer/scheduler/checkpointing intact.

Validation

Local:

• uv run pytest tests/test_train_cli.py -q -> 30 passed
• uv run ruff check src/benchflow/training/backends/prime_rl.py src/benchflow/cli/train.py tests/test_train_cli.py
• uv run ruff format --check src/benchflow/training/backends/prime_rl.py src/benchflow/cli/train.py tests/test_train_cli.py
• uv run ty check src/benchflow/training/backends/prime_rl.py src/benchflow/cli/train.py tests/test_train_cli.py
• git diff --check

H100 remote smoke:

• remote focused tests: 30 passed
• dry-run command: bench train run sft --compat-profile env0-mobile300-pr828 --data /home/ubuntu/work/mobile300-fixed/datasets/local_mobile300 --prime-rl-dir /home/ubuntu/work/prime-rl --dry-run
• dry-run manifest: overall_status=succeeded, custom_trainer_pretokenized_rows=300, custom_trainer_pretokenized_trainable_tokens=937095, derived_max_steps=37, tool_defs_removed_rows=300
• staged schema check: 300 rows, no messages/tools/tool_defs, tensor keys only, sum_trainable=937095
• Prime-RL import smoke: generated sitecustomize.py enabled both sample-mean and pre-tokenized data shims; SFTDataset._process returned the staged row tensors unchanged; prime_rl.trainer.sft.train imported successfully under the sample-mean shim

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 7e53897198

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Name the guarded commit in the regression test

The repo's AGENTS.md says regression tests must name the PR/commit they guard. This newly added test exercises the shifted-label pretokenized compatibility fix, but its docstring only describes the behavior and does not name 1c0bf71 or a PR, so it violates that review rule and loses the traceability the convention is meant to provide.

Useful? React with 👍 / 👎.

[bingran-you]

Additional E2E validation after padding fix

Validated commit b83dd1c7f46b4e248fa3ae009922a574569f3b28 on a single H100 against the real Mobile300 data.

• BenchFlow dry-run succeeded with --compat-profile env0-mobile300-pr828 and staged 300 pretokenized rows.
  • custom_trainer_pretokenized_rows=300
  • custom_trainer_pretokenized_trainable_tokens=937095
  • message_tail_truncated_rows=97
  • derived_max_steps=37, effective_train_examples=296
• Full Prime-RL SFT completed successfully on H100.
  • W&B: https://wandb.ai/benchflow-ai/env0-mobile-pr828-qwen35-prime-rl-sft-pretokenized-paddingfix-20260630/runs/mobile300_mobile300_pretokenized_pr866_paddingfix_prime_rl_main_20260630T163314Z
  • Checkpoint: step_37/STABLE
  • Loss: first 4.8484, final 0.4101, min 0.3809, max 4.8484 across 37 logged steps.
• Fireworks PEFT deployment smoke passed.
  • Plain chat returned HTTP 200.
  • Forced native tool-call request returned HTTP 200 with structured tool_calls and finish_reason=tool_calls.
• Daytona/OpenHands canary against 3 previously regressed tasks completed without infra errors.
  • Score: 0/3, errors 0, verifier errors 0, idle timeouts 0.
  • Tool calls: 35 total, 11.67 average, 0/3 zero-tool rows.
  • Per task: auth-service-account-impersonation 0.0 / 21 tools, gcal-ietf-interim-cancelled-sessions 0.0 / 11 tools, slack-do-not-kick-innocent 0.0 / 3 tools.

Conclusion: this PR fixes the BenchFlow-to-Prime-RL staging/trainer compatibility issue and removes the zero-tool collapse symptom seen in the previous sample-mean/token-suffix canary. It does not by itself recover pass rate on the 3-task behavioral canary, so the remaining regression should be investigated as model behavior/eval-runtime parity rather than a results.jsonl schema/load failure.