feat: support exclusive full-parameter training (train_mode: full) - #257
Merged
Conversation
- Server: add train_mode (lora|full) to ModelArgs with single-replica
validation; exclusive tenant lock; reload base weights on tenant release
- Backends: split TwinkleCompatFull{Transformers,Megatron}Model wrappers
(no PEFT/MultiLora) sharing the tinker/twinkle compat mixin
- Handlers: route tenant adapter names to the default optimizer group via
resolve_model_adapter_name; validate lora_config against train_mode
- Sampler: load full HF checkpoints into vLLM base model
(load_full_weights_from_path), dispatch on adapter_config.json presence
- Client: ServiceClient.create_full_training_client(_async) via patch_tinker
- Fix Megatron tinker_step ignoring AdamParams lr (param_groups not updated)
- Fix Megatron zero_grad/reload never zeroing DDP grad buffers
(hasattr on model chunk list was always False)
- Fix MultiLora.patch() crash on Megatron model chunk lists (#236 regression)
- Tests: full-param E2E (test_full_param_e2e.py); config variant presets in
start_e2e_server.py replacing per-combination yaml copies
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
tastelikefeet
approved these changes
Aug 14, 2026
There was a problem hiding this comment.
Pull request overview
Adds an exclusive full-parameter training mode (train_mode: full) to Twinkle’s unified server, enabling training of base model weights directly (no PEFT/MultiLoRA wrapping) with single-tenant exclusivity, plus end-to-end coverage and several backend/sampler fixes to support it.
Changes:
- Introduces
train_mode: fullacross server config, model management, and handlers, with exclusive-tenant enforcement and base-weight reset on tenant release. - Adds full-checkpoint load paths for backends and sampler (vLLM), and splits backend wrappers into shared compat mixins + LoRA/full concrete classes.
- Adds GPU E2E integration test for full-parameter training and enhances E2E server launcher with config variants and ad-hoc overrides.
Reviewed changes
Copilot reviewed 24 out of 25 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/server/start_e2e_server.py | Adds --variant presets and --set overrides to generate E2E server configs. |
| tests/server/model/test_tinker_handlers.py | Updates dummy management to include new full-mode fields/methods used by handlers. |
| tests/server/integration/test_full_param_e2e.py | New GPU E2E script validating full-parameter training, exclusivity guards, save + sample, and checkpoint format. |
| src/twinkle/server/sampler/twinkle_handlers.py | Allows adapter_uri to resolve to either LoRA adapter dir or full HF checkpoint, loading full weights into sampler when applicable. |
| src/twinkle/server/sampler/tinker_handlers.py | Same LoRA-vs-full checkpoint dispatch for tinker-compatible sampler endpoint. |
| src/twinkle/server/sampler/backends/mock_sampler.py | Adds a no-op load_full_weights_from_path for mock sampler backend. |
| src/twinkle/server/model/twinkle_handlers.py | Routes adapter_name through resolve_model_adapter_name, validates config vs train_mode, and enforces full-mode exclusivity (409). |
| src/twinkle/server/model/tinker_handlers.py | Adds full-mode validation/exclusive guard and routes training ops through resolved adapter name ('' in full mode). |
| src/twinkle/server/model/backends/transformers_model.py | Extracts compat mixin and adds full-mode wrapper around plain TransformersModel. |
| src/twinkle/server/model/backends/megatron_model.py | Extracts compat mixin and adds full-mode wrapper around plain MegatronModel; fixes Megatron optimizer param application. |
| src/twinkle/server/model/backends/mock_model.py | Adds no-op reload_initial_weights for mock backend. |
| src/twinkle/server/model/app.py | Adds train_mode plumbing, backend selection for LoRA vs full wrappers, exclusivity helpers, and reload-on-release in full mode. |
| src/twinkle/server/exceptions.py | Introduces FullModeBusyError for exclusive full-parameter deployments. |
| src/twinkle/server/config/application_spec.py | Adds ModelArgs.train_mode and config-load validation requiring single replica for full mode. |
| src/twinkle/sampler/vllm_sampler/vllm_sampler.py | Implements full-checkpoint streaming load into vLLM engine base weights (idempotent by resolved path). |
| src/twinkle/model/transformers/transformers.py | Implements non-LoRA load() branch (full HF checkpoint) and adds reload_initial_weights() for full-mode tenant reset; improves optimizer wrapper handling in step(). |
| src/twinkle/model/transformers/strategy/native_fsdp.py | Adds load_full_state_dict() to support full-checkpoint reload/load (incl. FSDP2 path). |
| src/twinkle/model/transformers/strategy/accelerate.py | Adds load_full_state_dict() to support full-checkpoint reload/load (incl. FSDP2 path). |
| src/twinkle/model/multi_lora.py | Fixes MultiLora.patch() device probing for Megatron chunk-list models. |
| src/twinkle/model/multi_lora_target_parameters.py | Makes target-parameters state-dict save/load tolerant for tenants without slots. |
| src/twinkle/model/megatron/megatron.py | Fixes Megatron DDP zero_grad_buffer for chunk-list models and adds reload_initial_weights() for full-mode tenant reset. |
| src/twinkle_client/utils/patch_tinker.py | Adds ServiceClient.create_full_training_client(_async) patch to request full-parameter mode (lora_config=None). |
| src/twinkle_client/types/model.py | Makes AddAdapterRequest.config optional (None for full-parameter). |
| src/twinkle_client/model/multi_lora_transformers.py | Allows add_adapter_to_model(..., config=None) for full-parameter deployments. |
| .gitignore | Ignores generated E2E config file. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+185
to
+188
| for dep in self.deployments: | ||
| num_replicas = dep.get('num_replicas') | ||
| max_replicas = (dep.get('autoscaling_config') or {}).get('max_replicas') | ||
| if (num_replicas or 1) > 1 or (max_replicas or 1) > 1: |
Comment on lines
+517
to
+519
| else: | ||
| files = sorted(glob.glob(os.path.join(resolved, '*.safetensors'))) | ||
| for fp in files: |
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
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.
Support exclusive full-parameter training (
train_mode: full)Summary
Adds an exclusive full-parameter training mode to the Twinkle server. A model deployment configured with
train_mode: fulltrains the base model weights directly (no PEFT/MultiLora wrapping), held by one tenant at a time. The tinker-compatible client gainscreate_full_training_client(), with a training loop identical to the LoRA path.Also fixes three pre-existing bugs surfaced by the new E2E coverage (see Bug fixes).
Key changes
Server
ModelArgs.train_mode: Literal['lora','full'](defaultlora), with a config-load-time validator forcing single replica in full mode (the exclusive-tenant lock lives in per-replica memory)ModelManagement:is_full_mode/resolve_model_adapter_name()(full mode routes every request to the default''optimizer group) /assert_full_mode_available()(raisesFullModeBusyErrorbefore any state registration)reload_initial_weights()— restores pristine base weights, zeroes DDP grad buffers andparam.grad, so the next tenant never inherits trained weightslora_configagainst the deployment mode (400/User on mismatch, 409 on exclusive-busy)Backends
TwinkleCompatFullTransformersModel/TwinkleCompatFullMegatronModel: wrap the plain models; the shared tinker/twinkle compat methods are extracted into mixins (method bodies only callsuper(), so they work against either MRO)TransformersModel.load()non-LoRA branch implemented (wasNotImplementedError): reads a full HF checkpoint (single/sharded safetensors or bin) and loads it in-place via the newstrategy.load_full_state_dict()(FSDP2 usesset_model_state_dict, plain models useload_state_dict)Sampler
vLLMSampler.load_full_weights_from_path(): streams a full HF checkpoint into the engine's base model (lazy per-tensor iteration, idempotent by resolved path), invalidates synced LoRA caches and resets the prefix cacheadapter_config.jsonpresence: LoRA adapter dir →adapter_path, full checkpoint → base-weight loadClient
ServiceClient.create_full_training_client/create_full_training_client_async(sendslora_config=None), returning the standardTrainingClientAddAdapterRequest.configbecomes optional (None= full-parameter)Bug fixes
tinker_stepignoredAdamParams.learning_rate— it only updatedchained_opt.config.lr, but Megatron readsparam_group['lr']at step time, so training silently ran at the construction-time default1e-4. Now passesoptim_paramstostep(), mirroring the transformers backend. (Affected the LoRA tinker path too.)hasattr(self.model, 'zero_grad_buffer')was alwaysFalsebecauseself.modelis aList[nn.Module]of chunks; gradients accumulated across iterations inmain_gradbuffers. Now iterates chunks. Together with (1) this caused loss spikes (2.7 → 7.5) in full-parameter Megatron training; after the fix the loss curve matches the transformers backend point-for-point.MultiLora.patch()crashed on Megatron chunk lists ('list' object has no attribute 'parameters', regression from Support multi-LoRA training with EP + FSDP2 #236) — Megatron LoRA server deployments could not start. Now probes the first chunk.Tests
tests/server/integration/test_full_param_e2e.py(5 phases: SFT loss decrease, LoRA-create rejected, second full tenant rejected, save+sample on full weights, on-disk checkpoint is a full HF checkpoint)start_e2e_server.pygains--variantpresets +--set key.path=valueoverrides that generate configs from the two base yamls (replaces per-combination config copies; generated configs verified dict-equal to the handwritten ones they replace)Verification (all on Qwen3.5-4B, real GPUs)
tests/server/model,tests/server/config)step-1 loss is consistent across all five parallel configs (2.684–2.689), and megatron/transformers full-parameter loss curves now match.
Known limitations / follow-ups
create_full_training_client(seed=...)is accepted but currently unused (kept for signature parity with the SDK)