[tml] Inkling native LoRA support - #2122
Merged
Merged
Conversation
Zhichenzzz
requested review from
Shi-Dong,
fzyzcjy,
guapisolo,
jybsuper,
maocheng23,
yueming-yuan and
yushengsu-thu
as code owners
August 3, 2026 20:10
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
yueming-yuan
approved these changes
Aug 3, 2026
Zhichenzzz
force-pushed
the
zhichen/feat-lora
branch
from
August 3, 2026 20:46
e707631 to
4354783
Compare
Adapter-only GRPO on the same backend and parallel stack as full-parameter training. The adapter follows Inkling's released LoRA schema (attention, dense MLP, shared-outer routed experts, shared experts, lm head; r=32 all-linear by default), applied per-module by the plugin at model build. After each step the exporter assembles the serving-ready adapter from the distributed training state (one flat all_gather per TP/EP group, PP broadcast) and hands it to the colocated SGLang engine over CUDA IPC; the frozen base is never re-transferred. Includes --lora-train-only (train the adapter, roll out the frozen base), warm starts from released safetensors adapters via --lora-adapter-path, the launcher's lora train-mode, and the 4-layer LoRA CI test with engine-side sha256 verification.
Zhichenzzz
force-pushed
the
zhichen/feat-lora
branch
from
August 3, 2026 21:00
4354783 to
2945d3d
Compare
Zhichenzzz
added a commit
that referenced
this pull request
Aug 3, 2026
Brings the Inkling model family and its native LoRA (#2122) onto the plugin branch. Resolution notes: - model.py keeps the generic resolve_lora_provider path; the 'inkling in custom_model_provider_path' special-cases are replaced by provider dispatch: resolve_lora_provider now routes Inkling checkpoints to miles_plugins.models.inkling.lora (which gains the provider-protocol aliases), and uses_builtin_native_lora_provider agrees with that dispatch. - sglang_engine keeps the plugin-normalized target pathway and adopts main's ['all'] sentinel for Inkling plus the adapter startup-preload behavior. - the --lora-train-only rename from main is reverted to --debug-lora-train-only per review; main's loadable-checkpoint guard in model.py is kept. - lora_utils keeps the native dist-ckpt save/load; main's rank-shard renames applied to the pre-refactor format and are superseded.
This was referenced Aug 3, 2026
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
2 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.
Stacked on #1683 (Inkling model + full-parameter RL); this PR adds the LoRA layer on top.
Adapter-only GRPO on the same backend and parallel stack: per-module adapters following Inkling's released LoRA schema (attention, dense MLP, shared-outer routed experts, shared experts, lm head; r=32 all-linear by default). After each step the exporter assembles a serving-ready adapter from the distributed training state (one flat all_gather per TP/EP group + PP broadcast) and hands it to the colocated SGLang engine over CUDA IPC — the frozen base is never re-transferred (weight update 49.4s → 2.5s, train step ~85% of full-parameter). Includes
--lora-train-only, warm starts via--lora-adapter-path, the launcher's--train-mode lora, and the 4-layer LoRA CI test with engine-side sha256 verification.Validated: 4-node Inkling-Small LoRA runs with steadily rising reward and aime25 (wandb
5kb81l5d,hxrf9erh); 4-layer CI green.https://www.lmsys.org/blog/2026-07-15-inkling-day0-support