[tml] Inkling model support - #1683
Merged
Merged
Conversation
Contributor
|
Note Gemini is unable to generate a review for this pull request due to the file types involved not being currently supported. |
Zhichenzzz
marked this pull request as ready for review
July 15, 2026 18:27
Zhichenzzz
requested review from
Shi-Dong,
fzyzcjy,
guapisolo,
jybsuper,
maocheng23,
yueming-yuan and
yushengsu-thu
as code owners
July 15, 2026 18:27
13 conflicting files. Where main had refactored or generalized the same code,
main wins; where the two sides added independent features, both are kept.
Take main:
- megatron_utils/arguments.py, sglang_utils/arguments.py: main is a strict
superset (multi-LoRA plain-DDP branch; the --sglang-* aliases 1683 guarded
with hasattr are gone, and --sglang-ep-size is now the real flag).
- ray/actor_group.py: main moved the body into allocate_gpus_for_actor().
Verified that helper carries everything 1683 added inline -- the TMS preload
path via get_binary_path_from_package, NVSHMEM_DISABLE_NCCL, and the disk
backup env vars.
- megatron_utils/actor.py reclaim: TMS_DISK_BACKUP_DIR is per-rank
(actor_factory: cell{i}_rank{r}), so main's wholesale rmtree of it is both
correct and simpler than globbing tms_*.bin, and hunk 4 must read that env
var rather than the shared root.
- utils/arguments.py: uid-scoped offload dir, the chunk-mb positivity assert,
sglang_cuda_graph_backend_prefill replacing the piecewise-cuda-graph flags
(nothing else on this branch referenced them), and the
async_max_concurrent_samples assert.
Keep both:
- sglang_engine.py: 1683's full-adapter serialized_tensors transport with its
exactly-one check AND main's in-place upsert; the Inkling "all" target-module
sentinel survives.
- update_weight_from_tensor.py: skip_base_sync now covers all three base
retention paths (distributed, cpu_backup, and 1683's colocate_base_persistent)
AND main's opt-out when --check-weight-update-equal is on. Main's
placeholder-rank guard is kept, but NOT its load_format="flattened_bucket":
this branch serializes a LoRA adapter as a plain named dict, so declaring a
flattened bucket would misread the payload. Main's to-do-1 described its own
bucket path and no longer applies; its CI-coverage to-do is kept.
- model.py: main's multi-LoRA hide_adapters load context composed with 1683's
missing-checkpoint tolerance and post-load native adapter load.
- lora_utils.py save: 1683 saves one shard per global rank because ranks
sharing (tp,pp) hold different EP experts, and the already-merged load side
reads that name with a legacy tp/pp fallback. Main's parallel_state
accessors are adopted so the HF PEFT write gate can use its pp_rank fix.
- The rollout-side files take main's module layout (LORA_ADAPTER_NAME and
friends now live in miles/utils/lora, not the megatron backend) while their
single-adapter branches gate on lora_rollout_enabled so --lora-train-only
still keeps the rollout side adapter-free.
Two breakages git could not see, both in non-conflicting regions:
- main renamed is_megatron_main_rank -> is_first_replica_megatron_main_rank;
model.py had two surviving calls to the old name.
- both sides registered --check-lora-weight-equal, which makes argparse raise
on startup. Main's is kept (identical behaviour, more precise help).
lora_rollout_enabled now lives in miles/utils/lora.py next to is_lora_enabled
so rollout code need not import the megatron backend, and is re-exported from
lora_utils.py for the training-side callers.
run_inkling_975b.py: the optimizer_args concatenation no longer needs to be split across lines after the NVMe args were dropped, and the comment above it still described them. run_inkling_975b_fully_async.py: a stray blank line after the docstring-only callback; pre-existing on this branch and failing the pinned black 24.3.0.
Three places where 1683's code and main's tests disagreed after the merge. git could not see any of them: none are in a conflicting region. 1. sglang_lora_target_all_sentinel read args.hf_checkpoint directly, so main's TestBuildLoraSyncConfig cases (a minimal Namespace) raised AttributeError. is_inkling_checkpoint already treats an empty path as 'not Inkling'. 2. setup_model_and_optimizer raised NotImplementedError for LoRA in raw mode unless the provider was the Inkling plugin, which broke main's test_lora_raw_mode_skips_bridge and changed behaviour for models this PR is not about. Apply the Inkling adapter wrapping when the Inkling provider is in use and otherwise leave main's path alone, so the PR stays additive. 3. get_named_update_units: 1683 replaced main's 'group references no params' assert with an unconditional continue -- a dead loop -- because a PP rank can legitimately hold no layer for an MoE-only group. Express that per group instead: AtomicUpdateGroup gains optional=False, Inkling's three MoE groups set it, and every other group keeps main's guard against a mistyped suffix.
…-train The bridge LoRA path patches _ParamAndGradBuffer so adapter param/grad buffers land in the "param_buffer" torch_memory_saver region, which pause(tag="default") leaves mapped. The native Inkling path skipped that patch, so under --offload-train the adapter buffers sat in the pausable region: sleep() unmapped them, and the by-design asleep update_weights hit an illegal memory access when the expert-parallel all_gather read the adapter weights. Production never saw this because engines with >=16 GPUs auto-disable offload; the 4-GPU CI topology was the first to combine LoRA with offload-train. Mirror the bridge gating in apply_inkling_lora.
_w1_delta/_a_cat are sglang-side buffers recomputed from the adapter at load time; the checker's pre-sync snapshot holds their uninitialized storage and can never match. Skip them like the frozen towers.
Zhichenzzz
force-pushed
the
zhichen/feat
branch
from
August 2, 2026 03:59
4d06f02 to
4f1086e
Compare
Same logic, new shape: module-level imports and adapter class, per-module appliers (attention / dense MLP / routed experts / shared experts / lm head), per-kind loader and exporter tables, and descriptive names throughout. The flat one-all_gather-per-group export batching is kept. Drops the unused megatron<->HF lora name mappers and the _APPLIED_MODELS registry.
This comment was marked as duplicate.
This comment was marked as duplicate.
Sorry, something went wrong.
yushengsu-thu
requested changes
Aug 2, 2026
load_inkling_lora_adapter writes the adapter params after load_checkpoint has already populated the distributed optimizer's fp32 masters, so the first step() would restore the pre-load init values over the adapter. Refresh the masters from the just-loaded model params.
…izer refactors The launcher predates two main-side moves: #1716 replaced the examples/fully_async module with miles.rollout.fully_async_rollout (selected by --fully-async + MILES_EXPERIMENTAL_ROLLOUT_REFACTOR=1), and the NVMe optimizer streaming merged as --stream-optimizer-state-to-disk / --offload-train-disk-dir rather than --optimizer-state-nvme-dir. Point the script at the current interfaces and drop the stale knob from run_inkling.
The gate still read args.inkling_mm_towers, but that CLI knob was removed when the plugin knobs folded into the model config: the towers are now chosen by the provider (inkling_mm_model_provider). The dead gate meant tower tensors were never appended to base syncs, leaving the engine with scrambled towers after its first offload cycle on multimodal runs.
wire_mm_towers runs with train=False, so the no_grad context covered not just the frozen tower forwards but also _scatter, whose clone of the embedding output detached it from the graph: embeddings received zero gradient on every multimodal batch. Scope the context to the towers.
The hand-rolled python renderer covered plain user/system text but not the rest of the chat spec: no thinking-effort line, no assistant tool_calls, no end-sampling token, and image/audio parts were rejected. inkling_fixed.jinja already implements all of it (and is what the TITO path serves), so route apply_chat_template through the bundled template and retire the python renderer.
The converter auto-bumps PP toward the rank count when launched with PP1 across many ranks, which is incompatible with the documented TP4 16-rank command and fails the world-size divisibility check.
yueming-yuan
reviewed
Aug 3, 2026
| else: | ||
| # Start async all_gather | ||
| if ".experts." in info.name: | ||
| if ".experts." in info.name and ".shared_experts." not in info.name: |
Collaborator
There was a problem hiding this comment.
".experts." seems to have more callings across many places... better to change all by adding not ".shared_experts."
yueming-yuan
reviewed
Aug 3, 2026
yueming-yuan
reviewed
Aug 3, 2026
yueming-yuan
reviewed
Aug 3, 2026
Collaborator
|
have some thoughts on lora checkpoint; will discuss offline |
Inkling's shared experts nest an inner ModuleList (mlp.shared_experts.experts.N), so a bare ".experts." substring test misclassifies their regular-TP params as expert-parallel. Three call sites already carried the ".shared_experts." exclusion inline; extract it as is_routed_expert_param() and use it at every megatron-name site (gather group choice, EP broadcast, full-size accounting, expert/dense iteration split). A no-op for models whose shared experts do not nest.
One launcher, one flag: --fully-async (with --rollout-num-nodes) selects the disaggregated topology - train_async.py driver, --fully-async rollout, bf16 grad reduce, NVMe-streamed optimizer state, and the validated TP4 PP2 EP16 profile for 8x4 training GPUs. Colocate arg assembly is unchanged. Also drop the two thin run-inkling-small*.sh wrappers; their validated recipes (full 5e-5 / lora 2e-4 at 64x128) now live in the doc launch commands.
…tions inkling.sh and inkling-small.sh already take MODEL_ARGS_NUM_LAYERS; the 4-layer files were two-line wrappers setting it, so let the launcher set it for *-4layer model names and point the registry at the base files. inkling-lora.sh defined a LORA_ARGS array nothing consumes (the launcher assembles its own LoRA flags) - delete it.
yushengsu-thu
approved these changes
Aug 3, 2026
This PR now carries the Inkling model + full-parameter RL only: the model plugin, launcher, multimodal path, chat template, weight-sync infra, docs, and the full-parameter 4-layer CI test. The LoRA adapter plugin, its launcher mode, engine serving wiring, and the LoRA CI test move to the follow-up PR stacked on this branch.
yueming-yuan
requested changes
Aug 3, 2026
The shim papered over dev-lineage sglang builds that lack begin/end_weight_update; the paired sglang here serves both, and swallowing 404s would mask a genuinely broken endpoint.
yueming-yuan
approved these changes
Aug 3, 2026
Zhichenzzz
added a commit
that referenced
this pull request
Aug 11, 2026
…nkling Kimi-K3 now has a page. It documents the day-0 LoRA RL recipe from #1825: the MXFP4 to BF16 upcast, the 32-rank torch_dist conversion, and the validated 16-node launch, with the LoRA target modules and what a healthy run looks like on the GB300 validation cluster. A line at the top says the branch it comes from, since none of it is on main yet. The page leans on two things worth calling out. LoRA is not a nicety here, it is what makes the recipe fit: the base stays frozen and only the adapters carry gradients, with one shared A factor and per-expert B factors across the 896 experts. And the memory figure in "what a healthy run looks like" is the one to watch, because an adapter export that leaks shows up as allocated memory not returning to baseline after a weight sync, long before the run dies. Inkling still said its implementation "is open at" #1683. That merged on 2026-08-03 and the code is on main, so the line is gone; the page is the documentation now. Both tables link Kimi-K3 at the page instead of the pull request. MiMo and GPT-OSS move below JoyAI so the families the home page lists appear in the same order in both, with the two it does not carry after them.
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
3 participants
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.
Inkling model + full-parameter RL support. The native LoRA layer is stacked on top in #2122.
https://www.lmsys.org/blog/2026-07-15-inkling-day0-support
sgl-project/sglang#31681
radixark/Megatron-LM#68