Feat/multi gpu training

[zhiheng-yang]

feat(training): add single-node multi-GPU SFT and GRPO

Summary

This PR adds single-node DDP and FSDP2 training to the existing SFT and GRPO pipeline, targeting bf16 LoRA on nvidia gpus.
Each training stage can select an Accelerate YAML profile; omitting the profile keeps the existing in-process launch path and CLI commands.

The local implementation has completed real gpu training checks for both SFT and online GRPO, including checkpoint export.

Launch and checkpoint handling

SFT and GRPO use the same launcher, rank helpers, and checkpoint export code.
Workers check their process count against the selected profile, and launch subprocesses receive an explicit environment so distributed-training variables do not leak into sandbox commands.

Profiles are tracked by content for resume checks and bundled with HF Jobs and prepared submissions.
Older run plans without a launch section remain compatible with the default in-process configuration.
Each FSDP2 profile must request a full state dict so saving produces a portable adapter; the coordinator exports the merged model on CPU after the training workers exit.

Scheduler allocations and site-specific cache paths stay outside the training code.
Training ranks, inference GPUs, and sandbox concurrency are configured separately.

Online GRPO

Each rollout has a unique global ID, and the configured sandbox concurrency is shared across ranks.
Verifier rewards are gathered in rollout order before calculating group advantages, so adding ranks does not change which samples belong to a reward group.

OpenCode can rewrite tool-call history, refresh context, or delegate to another agent.
That makes one flattened transcript an unreliable representation of every sampled action.
The collector therefore saves each model call's exact prompt IDs, completion IDs, and sampled logprobs, then trains those calls as separate causal segments:

real OpenCode rollouts -> verifier rewards -> rollout-group advantages
           |
           v
exact per-call token records -> aligned segment chunks -> policy update
                                                            |
                                                            v
                                               sync policy for next batch

Each segment keeps its original rollout's advantage.
Sequence-level importance correction is calculated across all segments of that rollout, and DAPO uses the global sampled-token count for normalization.
Old-policy probabilities are cached before optimization and reused across the corresponding training steps.

Long trajectories are processed one segment at a time, with backward immediately releasing each segment's graph.
Ranks agree on the number of execution slots and use zero-loss dummy slots when needed, keeping FSDP2 collectives aligned without dropping real actions.
This addresses the memory failure seen when all long segments were scored in one batch.

GRPO explicitly loads the pinned causal policy, and the model bridge preserves schema-declared numeric and structured XML tool arguments.
An optional OpenCode step budget applies to the built-in main and subagents; it remains a per-agent soft limit, separate from the overall timeout.

Main files

Python filenames below are under pipelines/benchflow-task-posttrain/src/posttrainarena/benchflow_pipeline/.

• config.py, launcher.py, cli.py, and pipeline.py: stage profiles, worker dispatch, timeout settings, and resume checks.
• distributed.py, sft.py, and checkpoint.py: rank coordination, SFT execution, adapter saving, and CPU merged export.
• grpo.py, segmented_rollout.py, and segmented_trainer.py: distributed collection, exact sampling records, and memory-bounded segment training.
• model_bridge.py and opencode.py: tool-argument handling and agent configuration.
• hf_jobs.py and submission.py: profile packaging for remote jobs and prepared submissions.
• The pipeline's configs/accelerate/ and tests/ directories: launch examples and regression coverage.

Validation

The full pipeline suite passed on the validated local tree: 362 passed.
Regression coverage includes context changes, exact segment preservation, rollout-group mapping, and cached old-policy probabilities; a two-process CPU test also checks unequal segment counts, collective ordering, and gradient equivalence.

SFT checks on real hpc GPUs covered single-process, DDP, and FSDP2 training on small supervised datasets, adapter saving, CPU merged export, and checkpoint reload.

The final online GRPO run used three GPUs: two training ranks and a separate inference GPU.
It ran fresh OpenCode tasks in Apptainer sandboxes, with real verifier rewards and no synthetic rollouts or substituted rewards.

Final online check	DDP	FSDP2
Real rollouts	8	8
Optimizer steps	4	4
Groups with reward variance	1	1
Nonzero saved LoRA-B tensors	248	248
Exact sampled tokens independently checked	39,256	34,780
Adapter and merged-model export checks	Passed	Passed

Both modes finished without agent or verifier errors, produced finite training results, and passed the saved-parameter and export-digest checks.
Every successful agent exchange was matched against its persisted token and logprob record.

The final multi-rank online GRPO checks used a local BenchFlow checkout with Apptainer support so agent sandboxes could run on the HPC allocation.
Apptainer is not a dependency of the training implementation; it was the sandbox backend used for this validation.
Related backend work is tracked in benchflow-ai/benchflow#1073; reproducing this run requires the locally modified checkout, not necessarily that PR alone.
Equivalent multi-rank online validation with the unchanged pinned upstream BenchFlow dependency and Daytona remains outstanding.
Earlier Daytona attempts encountered provider-routing errors and service constraints, so we used Apptainer on the HPC allocation to complete the online validation.
The validated scope is single-node, text-only bf16 LoRA; multi-node training, full fine-tuning, quantization, and other GRPO loss variants are outside this validation.

Spend and external systems

The final successful online validation used three HPC GPUs for approximately 1 hour 44 minutes.
A preceding replay used two GPUs for approximately 32 minutes and trained on the unchanged recorded trajectories to verify the memory fix under DDP and FSDP2.
Earlier SFT GPU checks, failed diagnostic attempts, and Daytona Oracle checks also consumed resources; the durations above are not the total development spend.

These experiments used the local integration worktree before the implementation was assembled into this branch.
Model caches, checkpoints, and raw traces remain local and should not be included in the PR.
No hosted-model API was used for the final local-model GRPO run.