Tinker seam re-port: decoupled fb/step on current upstream (supersedes fork main) - #2
Merged
Merged
Conversation
…er step) Tinker's API pipelines N forward_backward calls then one optim_step; upstream couples fwd+bwd+step inside train_one_step. This adds the minimal seam (specs/005 design in tinker-nemorl), reusing upstream data/loss/fanout: - model.py: extract _build_train_forward_step + _configure_train_step from train_one_step/train (pure moves); add forward_backward_pass (no step, grads accumulate across calls) and optimizer_step (per-call LR override; when the client sets LR the Megatron scheduler is not stepped). - loss: per-batch _loss_type_override (per-request loss selection) and _loss_norm_total (=1 -> pure-sum gradients, invariant to how a batch is split across forward_backward calls; fixes the 2x grad inflation measured in the G1 spike). Same rollout-key pattern as dynamic_global_batch_size. - data: forward the two scalar keys in get_batch; keep _partition_indices in process_rollout_data so per-sample outputs reassemble into client order. - actor: forward_backward_only / apply_optimizer_step / forward_logprobs / load_checkpoint (full resume incl. optimizer+scheduler). - ray/tinker_group.py (new): TinkerTrainGroup(RayTrainGroup) async fanout + DP-order merge. Frozen v1 RayTrainGroup untouched. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Tinker clients send advantages/log_probs/returns/values per sample (no actor-side compute_advantages_and_returns), and the seam's scalar keys (_loss_type_override, _loss_norm_total) must reach every rank. Add them to split_train_data_by_dp_raw's whitelists; absent keys are skipped, so non-Tinker paths are unaffected. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Tinker's fb contract includes per-datum logprobs (the SDK weights its chunk reduction by len(loss_fn_outputs) and the cookbook computes NLL from them). Computed via a forward-only pass on the same pre-step weights; extracting them from the loss pass itself is a follow-up optimization. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Fixes multi-fb grad accumulation on bridge-LoRA (G1 ratio 0.698 ~ 1/sqrt(2)). Probe evidence: fb2-entry buffer norms exactly equal fb1-exit (nothing zeroes), but fb2-exit barely moves / can DECREASE (219.23->219.35; 199.85->173.85) instead of matching the fresh fb(8) norms — the per-fb finalize (reduce-scatter) rewrites local grads with reduced shards, so the next backward accumulates onto corrupted contents. The raw full-FT path survived only by ddp_config luck (exact 1.0). Fix mirrors megatron's own microbatch accumulation: fb passes run with a no-op finalize_model_grads_func (restored after), and optimizer_step calls finalize_model_grads once over the locally-accumulated buffers before optimizer.step(). Runtime-verified 2026-07-14 (ns.config 4xH200, DP=4, Qwen2.5-0.5B bridge-LoRA): accumulation probe grad_norm bit-identical for 1xfb(8) vs 2xfb(4) + step (ratio exactly 1.0, was 0.698); full G1/G2 gates + sl_basic recipe rerun green. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
GavinZhu-GMI
force-pushed
the
tinker-seam
branch
from
July 14, 2026 01:21
40e2fad to
926f499
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Re-ports the Tinker orchestration interface onto current upstream radixark/miles (
fecd5d9), replacing the old fork implementation (main is 63 ahead / 648 behind upstream). Merging fast-forwards this fork to upstream + a minimal seam.The real review surface is the seam vs upstream: 6 commits, +392/−71 over 9 files (compare vs fecd5d9) — everything else in this PR's diff is upstream catching up.
The seam
forward_backward_pass/optimizer_stepsplit out oftrain_one_step(megatron_utils/model.py) — fwd/bwd with no optimizer step and no zero_grad, then a decoupled step with per-calllearning_ratereturning{success, grad_norm}.forward_backward_passswapsconfig.finalize_model_grads_functo a no-op (restored infinally);optimizer_steprunsfinalize_model_grads_with_empty_cacheonce beforeoptimizer.step()— mirroring Megatron's own microbatch accumulation. Also saves one DP reduce per fb call._loss_norm_total=1batch key — gradient is split-invariant across any fb pipelining depth; normalization stays in client weights per the Tinker contract._loss_type_overrideper-request loss selection;_partition_indicesfor client-order per-sample outputs; fb returns per-sample response log-probs (computed on pre-step weights).TinkerTrainGroupasync Ray fanout (miles/ray/tinker_group.py); DP-split passthrough for client-supplied advantages/log_probs;MILES_REPOdocker build arg.Validation (ns.config 4×H200, DP=4, Qwen2.5-0.5B bridge-LoRA, 2026-07-13/14)
0e89c58(marked UNVERIFIED in its commit message at push time): before the deferred finalization, the per-fb-call reduce-scatter corrupted the grad buffer as an accumulation substrate (ratio 0.698 on bridge-LoRA).Companion PR (backend adaptation): GMISWE/tinker-cloud#15.
Follow-up candidates: upstream the seam to radixark/miles; extend loss_hub if recipes need per-sample logprobs beyond the response slice;
load_checkpoint(path)full-resume on the actor.