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

.gitignore

@@ -158,3 +158,4 @@ ast_index_file.py
 test_cookbook/
 /test*.py
 swanlog/
+tests/server/config/_generated_e2e.yaml

src/twinkle/model/megatron/megatron.py

@@ -650,8 +650,11 @@ def zero_grad(self, **kwargs):
         # For DDP-wrapped models, ALWAYS zero the gradient buffer
         # This is essential because Megatron's forward_backward_func uses
         # the buffer's state to track gradient accumulation
-        if self._is_model_ddp_wrapped() and hasattr(self.model, 'zero_grad_buffer'):
-            self.model.zero_grad_buffer()
+        # (self.model is a list of chunks; zero_grad_buffer lives on each chunk)
+        if self._is_model_ddp_wrapped():
+            for model_chunk in self.model:
+                if hasattr(model_chunk, 'zero_grad_buffer'):
+                    model_chunk.zero_grad_buffer()
 
         if not optimizer_config.do_grad_sync(kwargs.pop('gradient_accumulation_steps', None)):
             return
@@ -988,6 +991,32 @@ def load(self, name: str, output_dir: Optional[str] = None, **kwargs):
         if dist.is_initialized():
             dist.barrier()
 
+    @remote_function(dispatch='all')
+    def reload_initial_weights(self, **kwargs):
+        """Reload the base model weights from the original model path.
+
+        Used by full-parameter server deployments after a tenant releases the
+        (exclusive) model, so the next tenant starts from clean pretrained
+        weights instead of the previous tenant's trained weights.
+        """
+        bridge = self.strategy.bridge
+        bridge.load_weights(
+            self.strategy.unwrap_model(self.model),
+            self._model_path,
+            peft_format=False,
+        )
+        # Drop any leftover gradients from the previous tenant so they cannot
+        # leak into the next tenant's first optimizer step.
+        if self._is_model_ddp_wrapped():
+            for model_chunk in self.model:
+                if hasattr(model_chunk, 'zero_grad_buffer'):
+                    model_chunk.zero_grad_buffer()
+        for _model in self.strategy.unwrap_model(self.model):
+            for param in _model.parameters():
+                param.grad = None
+        if dist.is_initialized():
+            dist.barrier()
+
     @remote_function(dispatch='all')
     def resume_from_checkpoint(self, checkpoint_dir, *, resume_only_model=False, **kwargs):
         adapter_name = kwargs.pop('adapter_name', self._get_default_group())

src/twinkle/model/multi_lora.py

@@ -479,9 +479,12 @@ def patch(self,
               target_modules='all-linear',
               *args,
               **kwargs):
-        module_device = getattr(module, 'device', None)
+        # ``module`` may be a list of model chunks (Megatron); probe the first
+        # chunk for the device in that case.
+        first_module = module[0] if isinstance(module, (list, tuple)) else module
+        module_device = getattr(first_module, 'device', None)
         if module_device is None:
-            module_device = next(module.parameters())[1].device
+            module_device = next(first_module.parameters()).device
         low_cpu_mem_usage = module_device.type == 'meta'
 
         for i in range(self.max_loras):

src/twinkle/model/multi_lora_target_parameters.py

@@ -369,14 +369,21 @@ def named_slot_parameters(self, tenant_adapter_name: str) -> Iterator[tuple[str,
             yield from wrapper.named_slot_parameters(slot_name)
 
     def get_state_dict(self, tenant_adapter_name: str) -> dict[str, torch.Tensor]:
-        slot_name = self.tenant_to_slot[tenant_adapter_name]
+        # Tenants without target_parameters never acquire a slot here; return
+        # an empty dict so plain-LoRA save paths do not fail.
+        slot_name = self.tenant_to_slot.get(tenant_adapter_name)
+        if slot_name is None:
+            return {}
         state_dict = {}
         for wrapper in self.wrappers:
             state_dict.update(wrapper.get_state_dict(slot_name))
         return state_dict
 
     def set_state_dict(self, tenant_adapter_name: str, state_dict: dict[str, torch.Tensor]) -> set[str]:
-        slot_name = self.tenant_to_slot[tenant_adapter_name]
+        # Same tolerance as get_state_dict for plain-LoRA load paths.
+        slot_name = self.tenant_to_slot.get(tenant_adapter_name)
+        if slot_name is None:
+            return set()
         consumed_keys = set()
         for wrapper in self.wrappers:
             consumed_keys.update(wrapper.set_state_dict(slot_name, state_dict))

src/twinkle/model/transformers/strategy/accelerate.py

@@ -203,6 +203,20 @@ def get_full_state_dict(self, model) -> dict:
             del local
         return state_dict
 
+    def load_full_state_dict(self, model, state_dict) -> None:
+        """Load a full (non-sharded) state dict into the model in-place.
+
+        Used by full-parameter training to (re)load base weights, e.g. after a
+        tenant releases an exclusive full-parameter deployment.
+        """
+        fsdp_plugin = self._get_fsdp_plugin()
+        if fsdp_plugin is not None and fsdp_plugin.fsdp_version == 2:
+            from torch.distributed.checkpoint.state_dict import set_model_state_dict
+            set_model_state_dict(model, state_dict, options=self._prepare_fsdp2_sd_options())
+            return
+        unwrapped = self.unwrap_model(model)
+        unwrapped.load_state_dict(state_dict, strict=False)
+
     def get_adapter_state_dict(self, model, adapter_name: str) -> dict:
         """Collect only LoRA adapter parameters."""
         from twinkle.utils import torch_util

src/twinkle/model/transformers/strategy/native_fsdp.py

@@ -345,6 +345,24 @@ def get_full_state_dict(self, model) -> dict:
 
         return state_dict
 
+    def load_full_state_dict(self, model, state_dict) -> None:
+        """Load a full (non-sharded) state dict into the (possibly sharded) model.
+
+        Used by full-parameter training to (re)load base weights. Uses
+        ``set_model_state_dict`` when the model is distributed (FSDP2), else a
+        plain in-place ``load_state_dict``.
+        """
+        if self.device_mesh is not None:
+            from torch.distributed.checkpoint.state_dict import StateDictOptions, set_model_state_dict
+            set_model_state_dict(
+                model,
+                state_dict,
+                options=StateDictOptions(full_state_dict=True, broadcast_from_rank0=True),
+            )
+            return
+        unwrapped = self.unwrap_model(model)
+        unwrapped.load_state_dict(state_dict, strict=False)
+
     def get_adapter_state_dict(self, model, adapter_name: str) -> dict:
         """Collect only LoRA adapter parameters, with EP-aware all-gather."""
         unwrapped = self.unwrap_model(model)

src/twinkle/model/transformers/transformers.py

@@ -146,6 +146,44 @@ def accumulate_metrics(self, is_training):
 DEFAULT_WEIGHT_DECAY = 0.01
 
 
+def _read_hf_state_dict(checkpoint_dir: str) -> Dict[str, torch.Tensor]:
+    """Read a full HuggingFace checkpoint directory into a CPU state dict.
+
+    Supports both single-file and sharded ``safetensors`` layouts, falling back
+    to ``pytorch_model.bin`` variants. Returns tensors on CPU.
+    """
+    import json
+
+    state_dict: Dict[str, torch.Tensor] = {}
+    st_index = os.path.join(checkpoint_dir, 'model.safetensors.index.json')
+    st_single = os.path.join(checkpoint_dir, 'model.safetensors')
+    bin_index = os.path.join(checkpoint_dir, 'pytorch_model.bin.index.json')
+    bin_single = os.path.join(checkpoint_dir, 'pytorch_model.bin')
+
+    if os.path.exists(st_index) or os.path.exists(st_single):
+        from safetensors.torch import load_file
+        if os.path.exists(st_index):
+            with open(st_index) as f:
+                weight_map = json.load(f)['weight_map']
+            shards = sorted(set(weight_map.values()))
+            for shard in shards:
+                state_dict.update(load_file(os.path.join(checkpoint_dir, shard), device='cpu'))
+        else:
+            state_dict.update(load_file(st_single, device='cpu'))
+    elif os.path.exists(bin_index) or os.path.exists(bin_single):
+        if os.path.exists(bin_index):
+            with open(bin_index) as f:
+                weight_map = json.load(f)['weight_map']
+            shards = sorted(set(weight_map.values()))
+            for shard in shards:
+                state_dict.update(torch.load(os.path.join(checkpoint_dir, shard), map_location='cpu', weights_only=True))
+        else:
+            state_dict.update(torch.load(bin_single, map_location='cpu', weights_only=True))
+    else:
+        raise FileNotFoundError(f'No safetensors/bin weights found in {checkpoint_dir}')
+    return state_dict
+
+
 @remote_class()
 class TransformersModel(TwinkleModel, PreTrainedModel, CheckpointEngineMixin):
     """The transformers model wrapper.
@@ -856,7 +894,11 @@ def step(self, **kwargs):
 
         optim_params = kwargs.pop('optim_params', {})
         if optim_params:
-            assert isinstance(optimizer, (AdamW, Adam))
+            # After _lazy_wrap_model the optimizer may be wrapped (e.g.
+            # accelerate's AcceleratedOptimizer); check the inner instance.
+            inner_optimizer = getattr(optimizer, 'optimizer', optimizer)
+            assert isinstance(inner_optimizer, (AdamW, Adam)), \
+                f'optim_params is only supported for Adam/AdamW, got {type(inner_optimizer).__name__}'
             for group in optimizer.param_groups:
                 group['lr'] = optim_params['lr']
                 if group['weight_decay'] > 0.0 and optim_params.get('weight_decay', None) is not None:
@@ -1137,11 +1179,31 @@ def load(self, name: str, output_dir: Optional[str] = None, **kwargs):
             adapter_weights = load_peft_weights(checkpoint_dir, device='cpu')
             self.strategy.load_peft_weights(model, adapter_weights, adapter_name)
         else:
-            raise NotImplementedError
+            # Full-parameter model: load a plain HF checkpoint in-place.
+            state_dict = _read_hf_state_dict(checkpoint_dir)
+            self.strategy.load_full_state_dict(self.model, state_dict)
 
         if load_optimizer:
             self._load_optimizer(checkpoint_dir, adapter_name=adapter_name)
 
+    @remote_function()
+    def reload_initial_weights(self, **kwargs):
+        """Reload the base model weights from ``self.model_id``.
+
+        Used by full-parameter server deployments after a tenant releases the
+        (exclusive) model, so the next tenant starts from clean pretrained
+        weights instead of the previous tenant's trained weights.
+        """
+        if not self.model_id:
+            logger.warning('reload_initial_weights skipped: model_id is not set (blank model).')
+            return
+        state_dict = _read_hf_state_dict(self.model_id)
+        self.strategy.load_full_state_dict(self.model, state_dict)
+        # Drop any leftover gradients from the previous tenant so they cannot
+        # leak into the next tenant's first optimizer step.
+        for param in self.model.parameters():
+            param.grad = None
+
     def _load_optimizer(self, checkpoint_dir, **kwargs):
         adapter_name = kwargs.pop('adapter_name', _default_adapter_name)
         strict = kwargs.pop('strict', False)

src/twinkle/sampler/vllm_sampler/vllm_sampler.py

@@ -486,6 +486,52 @@ async def _receive_and_load():
 
         self._run_in_loop(_receive_and_load())
 
+    @remote_function(dispatch='all', collect='first', lazy_collect=False)
+    def load_full_weights_from_path(self, path: str) -> int:
+        """Load a full (non-LoRA) HF checkpoint into the engine's base model.
+
+        Used by full-parameter training: the saved checkpoint is a plain HF
+        directory (no ``adapter_config.json``), so it replaces the sampler's
+        base weights instead of being loaded as a LoRA adapter. Idempotent:
+        repeated calls with the same resolved path are skipped.
+
+        Returns:
+            1 if weights were (re)loaded, 0 if the path was already loaded.
+        """
+        import glob
+        import json
+        import os
+
+        resolved = HubOperation.download_model(model_id_or_path=path)
+        if getattr(self, '_loaded_full_weights_path', None) == resolved:
+            return 0
+
+        from safetensors import safe_open
+
+        def _weight_iter():
+            index = os.path.join(resolved, 'model.safetensors.index.json')
+            if os.path.exists(index):
+                with open(index) as f:
+                    shards = sorted(set(json.load(f)['weight_map'].values()))
+                files = [os.path.join(resolved, s) for s in shards]
+            else:
+                files = sorted(glob.glob(os.path.join(resolved, '*.safetensors')))
+            for fp in files:
+                with safe_open(fp, framework='pt', device='cpu') as f:
+                    for key in f.keys():
+                        yield key, f.get_tensor(key)
+
+        async def _load():
+            await self.engine.update_weights(_weight_iter(), peft_config=None, base_sync_done=False)
+            # A full base-model swap invalidates any previously synced LoRA.
+            self.engine.invalidate_synced_lora()
+
+        logger.info(f'Loading full-parameter weights into sampler base model from {resolved}')
+        self._run_in_loop(_load())
+        self._loaded_full_weights_path = resolved
+        self.reset_prefix_cache()
+        return 1
+
     @remote_function(dispatch='all', collect='first', lazy_collect=False)
     def shutdown(self):
         """Gracefully shutdown the vLLM engine and background event loop.

src/twinkle/server/config/application_spec.py

@@ -54,6 +54,7 @@ class ModelArgs(_ArgsBase):
     device_group: dict[str, Any]
     device_mesh: dict[str, Any]
     backend: Literal['mock', 'transformers', 'megatron']
+    train_mode: Literal['lora', 'full'] = 'lora'
     adapter_config: dict[str, Any] | None = None
     queue_config: TaskQueueConfig = Field(default_factory=TaskQueueConfig)
     max_loras: int = 5
@@ -169,3 +170,24 @@ def _coerce_args_to_schema(cls, data: Any) -> Any:
         # ``schema.model_validate`` rejects a non-dict itself with a clean
         # error, so no separate non-dict guard is needed here.
         return {**data, 'args': schema.model_validate(raw_args)}
+
+    @model_validator(mode='after')
+    def _validate_full_mode_single_replica(self) -> ApplicationSpec:
+        """``train_mode: full`` requires exactly one replica.
+
+        The exclusive-tenant lock lives in per-replica memory
+        (``ModelManagement._resource_records``), so more than one replica would
+        silently allow one full-parameter tenant per replica, each rewriting
+        its own copy of the base weights. Reject that at config-load time.
+        """
+        if not (isinstance(self.args, ModelArgs) and self.args.train_mode == 'full'):
+            return self
+        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:
+                raise ValueError(
+                    f"Application '{self.name}': train_mode='full' is an exclusive single-tenant "
+                    'mode and requires a single replica; set num_replicas/autoscaling_config.'
+                    'max_replicas to 1.')
+        return self

src/twinkle/server/exceptions.py

@@ -51,3 +51,17 @@ class ConfigParseError(TwinkleServerError):
 class ResourceExhaustedError(TwinkleServerError):
     """Resource exhausted — queue full, insufficient memory, connection pool exhausted, etc."""
     pass
+
+
+class FullModeBusyError(TwinkleServerError):
+    """A full-parameter (exclusive) model deployment already has a holder.
+
+    Full-parameter training rewrites the shared base-model weights, so a single
+    deployment can only host one training task at a time.
+    """
+
+    def __init__(self, current_holder: str) -> None:
+        self.current_holder = current_holder
+        super().__init__('This deployment runs in full-parameter (exclusive) mode and is already '
+                         f'held by another training task ({current_holder}). Only one full-parameter '
+                         'training task is allowed at a time; retry after it is released.')