tinker seam: re-port onto upstream 2cb59235 (multi-LoRA era) + pin sglang-miles - #3
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>
Branch-tip (FETCH_HEAD) builds are unreproducible; pin the SHA that carries the multi-LoRA serving fixes (sglang#32376, #32421), matching the SHA-pin convention already used for Megatron-Bridge.
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.
Second seam re-port, same two-phase shape as PR #2: main has been hard-synced to upstream
2cb59235(130 commits, incl. the multi-LoRA stack radixark#1742-radixark#1748/radixark#1794 and sglang v0.5.16 bump); this PR rebases the 6-commit tinker seam on top and adds one pin commit.Conflicts resolved in the replay
training_utils/loss.py: upstream multi-LoRA and the seam both overrideglobal_batch_sizefor deferred normalization — both retained, seam's_loss_norm_totaltakes precedence when present (the two paths are disjoint today).megatron_utils/model.py: seam's_build_train_forward_stepextraction re-done around upstream's multi-LoRAtrain_one_step(retain-grads preamble + per-slot stepping preserved verbatim; the shared forward_step body keepsset_tokens_per_adapter_slot).utils/data.py/train_data_conversion.py:_partition_indicesfollows the partition logic intoprocess_rollout_data_shard; key-list unions.New commit
docker: pin SGLANG_COMMIT to sglang-miles tip 3003d70f6— FETCH_HEAD builds are unreproducible; the pinned SHA carries sglang#32376/#32421 (multi-LoRA serving fixes), matching the existing Megatron-Bridge SHA-pin convention.Old main archived at
main-archive-20260730(=df3f75cd, the 20260715 image source).Merge via rebase-merge to keep main linear (= upstream + seam + pin).