feat: support exclusive full-parameter training (train_mode: full)

[Yunnglin]

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: full trains the base model weights directly (no PEFT/MultiLora wrapping), held by one tenant at a time. The tinker-compatible client gains create_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'] (default lora), 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() (raises FullModeBusyError before any state registration)
• Tenant release in full mode calls reload_initial_weights() — restores pristine base weights, zeroes DDP grad buffers and param.grad, so the next tenant never inherits trained weights
• tinker/twinkle create routes validate lora_config against 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 call super(), so they work against either MRO)
• TransformersModel.load() non-LoRA branch implemented (was NotImplementedError): reads a full HF checkpoint (single/sharded safetensors or bin) and loads it in-place via the new strategy.load_full_state_dict() (FSDP2 uses set_model_state_dict, plain models use load_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 cache
• Sampler handlers dispatch on adapter_config.json presence: LoRA adapter dir → adapter_path, full checkpoint → base-weight load

Client

• ServiceClient.create_full_training_client / create_full_training_client_async (sends lora_config=None), returning the standard TrainingClient
• AddAdapterRequest.config becomes optional (None = full-parameter)

Bug fixes

1. Megatron tinker_step ignored AdamParams.learning_rate — it only updated chained_opt.config.lr, but Megatron reads param_group['lr'] at step time, so training silently ran at the construction-time default 1e-4. Now passes optim_params to step(), mirroring the transformers backend. (Affected the LoRA tinker path too.)
2. Megatron DDP grad buffers were never zeroed — hasattr(self.model, 'zero_grad_buffer') was always False because self.model is a List[nn.Module] of chunks; gradients accumulated across iterations in main_grad buffers. 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.
3. 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

• New 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.py gains --variant presets + --set key.path=value overrides 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)

Config	Result
transformers full, DP=2	ALL PHASES PASSED
transformers full, FSDP2	ALL PHASES PASSED
megatron full, DP=2 PP=2	ALL PHASES PASSED
megatron full, TP=2 PP=2	ALL PHASES PASSED
megatron full, TP=2 DP=2	ALL PHASES PASSED
megatron LoRA SFT regression (twinkle+tinker)	PASSED
transformers LoRA full-cycle regression	PASSED
unit tests (tests/server/model, tests/server/config)	21 passed

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

• Full mode is single-replica by design; the busy check has a narrow race window against the lifecycle expiry rollback path (documented, low likelihood) — follow-up candidate
• create_full_training_client(seed=...) is accepted but currently unused (kept for signature parity with the SDK)

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.