fix: preserve training state across checkpoints and shared adapters

src/twinkle/model/megatron/megatron.py

@@ -17,7 +17,7 @@
 from peft.tuners.lora import Linear as LoraLinear
 from torch.optim import Optimizer
 from torch.optim.lr_scheduler import LRScheduler
-from transformers import PreTrainedConfig
+from transformers import PretrainedConfig
 from typing import Any, Callable, Dict, Generator, List, Literal, Optional, Tuple, Type, Union
 
 import twinkle
@@ -101,7 +101,7 @@ class MegatronModel(TwinkleModel, nn.Module, CheckpointEngineMixin):
     def __init__(
         self,
         model_id: str,
-        config: Optional[PreTrainedConfig] = None,
+        config: Optional[PretrainedConfig] = None,
         ddp_config: Optional[Dict[str, Any]] = None,
         device_mesh: Optional[DeviceMesh] = None,
         mixed_precision: Literal['no', 'fp16', 'bf16'] = 'bf16',
@@ -980,6 +980,9 @@ def load(self, name: str, output_dir: Optional[str] = None, **kwargs):
                 checkpoint_dir,
                 peft_format=(adapter_name != _default_adapter_name),
             )
+            optimizer_config = self.optimizer_group.get(adapter_name)
+            if optimizer_config is not None and optimizer_config.optimizer is not None:
+                optimizer_config.optimizer.reload_model_params()
 
         if dist.is_initialized():
             dist.barrier()
@@ -1010,7 +1013,7 @@ def reload_initial_weights(self, **kwargs):
         if dist.is_initialized():
             dist.barrier()
 
-    @remote_function(dispatch='all')
+    @remote_function(dispatch='all', collect='first', sync=True, timeout=3600)
     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/megatron/multi_lora_megatron.py

@@ -365,6 +365,9 @@ def load(self, name: str, output_dir: Optional[str] = None, **kwargs):
                 adapter_name=real_adapter_name,
                 converter=self.multi_adapter.load_lora_converter)
 
+        optimizer_config = self.optimizer_group.get(adapter_name)
+        if optimizer_config is not None and optimizer_config.optimizer is not None:
+            optimizer_config.optimizer.reload_model_params()
         if load_optimizer:
             with self.optimizer_context(real_adapter_name):
                 self._load_multi_lora_optimizer(checkpoint_dir, adapter_name=adapter_name, **kwargs)

src/twinkle/model/megatron/strategy/megatron.py

@@ -1,7 +1,7 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
 import torch
 import torch.nn as nn
-from transformers import PreTrainedConfig
+from transformers import PretrainedConfig
 from typing import Any, Dict, List, Literal, Optional
 
 from twinkle import DeviceMesh, Platform, torch_util
@@ -44,7 +44,7 @@ def __init__(
         mixed_precision: Literal['no', 'fp16', 'bf16'] = 'bf16',
         seed: int = 42,
         variable_seq_lengths: bool = True,
-        config: PreTrainedConfig = None,
+        config: PretrainedConfig = None,
         ddp_config: Dict[str, Any] = None,
         **kwargs,
     ):
@@ -263,7 +263,7 @@ def reduce_loss(self, local_loss, local_count, logits, logps):
 
     def get_model_config(
         self,
-        hf_config: PreTrainedConfig,
+        hf_config: PretrainedConfig,
         parallel_kwargs: Dict[str, Any],
         **kwargs,
     ):

src/twinkle/model/multi_lora.py

@@ -1,6 +1,6 @@
 import re
 import torch
-from contextlib import contextmanager
+from contextlib import contextmanager, nullcontext
 from copy import deepcopy
 from dataclasses import dataclass, field
 from peft import LoraConfig, PeftModel, get_peft_model
@@ -413,8 +413,6 @@ def _embedding_forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
             def _megatron_forward(self, x: torch.Tensor, *args: Any, **kwargs: Any):
                 from megatron.core.extensions.transformer_engine import (TEGroupedLinear,
                                                                          TELayerNormColumnParallelLinear, TELinear)
-                from megatron.core.tensor_parallel import (gather_from_sequence_parallel_region,
-                                                           scatter_to_sequence_parallel_region)
                 from megatron.core.transformer.moe.router import TopKRouter
 
                 previous_dtype = x.dtype
@@ -440,9 +438,6 @@ def _megatron_forward(self, x: torch.Tensor, *args: Any, **kwargs: Any):
                     raise ValueError(f'Unsupported base layer type: {type(self.base_layer)}')
 
                 if not isinstance(self.base_layer, TopKRouter) and not self.disable_adapters and not self.merged:
-                    if self.sequence_parallel and self.base_layer.parallel_mode == 'column':
-                        x = gather_from_sequence_parallel_region(x)
-
                     for active_adapter in self.active_adapters:
                         if active_adapter not in self.lora_A.keys():
                             continue
@@ -458,56 +453,31 @@ def _megatron_forward(self, x: torch.Tensor, *args: Any, **kwargs: Any):
                         scaling = _lora.tenant_config.lora_alpha / (
                             _lora.tenant_config.r**0.5 if _lora.tenant_config.use_rslora else _lora.tenant_config.r)
 
-                        def _lora_A(x, *args, **kwargs):
-                            if isinstance(lora_A, TEGroupedLinear):
-
-                                def _get_weight_tensors(self):
-                                    tensors = self._get_weight_tensors_origin()
-                                    return [t[:_lora.tenant_config.r, :].contiguous() for t in tensors]
-
-                                lora_A._get_weight_tensors_origin = lora_A._get_weight_tensors
-                                lora_A._get_weight_tensors = MethodType(_get_weight_tensors, lora_A)
-                                output = lora_A(x, *args, **kwargs)
-                                lora_A._get_weight_tensors = lora_A._get_weight_tensors_origin
-                                delattr(lora_A, '_get_weight_tensors_origin')
-                                return output
-                            else:
-                                return torch.nn.functional.linear(
-                                    x, lora_A.weight[:_lora.tenant_config.r, :], bias=None)
-
-                        def _lora_B(x, *args, **kwargs):
-                            if isinstance(lora_B, TEGroupedLinear):
-
-                                def _get_weight_tensors(self):
-                                    tensors = self._get_weight_tensors_origin()
-                                    return [t[:, :_lora.tenant_config.r].contiguous() for t in tensors]
-
-                                lora_B._get_weight_tensors_origin = lora_B._get_weight_tensors
-                                lora_B._get_weight_tensors = MethodType(_get_weight_tensors, lora_B)
-                                output = lora_B(x, *args, **kwargs)
-                                lora_B._get_weight_tensors = lora_B._get_weight_tensors_origin
-                                delattr(lora_B, '_get_weight_tensors_origin')
-                                return output
-                            else:
-                                return torch.nn.functional.linear(
-                                    x, lora_B.weight[:, :_lora.tenant_config.r], bias=None)
-
                         dtype = lora_A.weight0.dtype if isinstance(lora_A, TEGroupedLinear) else lora_A.weight.dtype
                         x = x.to(dtype)
 
-                        lora_result = _lora_A(dropout(x), *args, **kwargs)
-                        if isinstance(lora_result, tuple):
-                            lora_result = lora_result[0]
-
-                        lora_result = _lora_B(lora_result, *args, **kwargs)
-                        if isinstance(lora_result, tuple):
-                            lora_result = lora_result[0]
-
+                        fp8_context = nullcontext()
+                        if getattr(self.config, 'fp8_param', False):
+                            import transformer_engine.pytorch as te
+                            fp8_context = te.fp8_autocast(enabled=False)
+                        with fp8_context:
+                            # Keep the backend layer's TP/SP communication and parameter
+                            # identities. Mask inactive ranks between A and B instead of
+                            # replacing parallel layers with local linear operations.
+                            lora_result = (
+                                lora_A(dropout(x), *args, **kwargs) if isinstance(lora_A, TEGroupedLinear) else lora_A(
+                                    dropout(x)))
+                            if isinstance(lora_result, tuple):
+                                lora_result = lora_result[0]
+                            lora_result = torch.nn.functional.pad(lora_result[..., :_lora.tenant_config.r],
+                                                                  (0, _self.max_r - _lora.tenant_config.r))
+                            lora_result = (
+                                lora_B(lora_result, *args, **kwargs)
+                                if isinstance(lora_B, TEGroupedLinear) else lora_B(lora_result))
+                            if isinstance(lora_result, tuple):
+                                lora_result = lora_result[0]
                         lora_result = lora_result * scaling
 
-                        if self.sequence_parallel and self.base_layer.parallel_mode == 'row':
-                            lora_result = scatter_to_sequence_parallel_region(lora_result)
-
                         result = result + lora_result
 
                 result = result.to(previous_dtype)

src/twinkle/model/transformers/multi_lora_transformers.py

@@ -311,8 +311,9 @@ def _optimizer_param_name_mapping(self, adapter_name: str, optimizer: Optimizer)
 
     @remote_function(collect='first')
     def save(self, name, output_dir: Optional[str] = None, interval=1, **kwargs):
-        self._check_adapter_valid(kwargs.get('adapter_name'))
-        with self.multi_adapter.save_context(kwargs.get('adapter_name')):
+        adapter_name = kwargs.get('adapter_name')
+        self._check_adapter_valid(adapter_name)
+        with self._adapter_context(adapter_name), self.multi_adapter.save_context(adapter_name):
             checkpoint_dir = super().save(name, output_dir, interval, **kwargs)
         if dist.is_initialized():
             dist.barrier()
@@ -322,7 +323,7 @@ def save(self, name, output_dir: Optional[str] = None, interval=1, **kwargs):
     def load(self, name: str, output_dir: Optional[str] = None, **kwargs):
         adapter_name = kwargs.get('adapter_name')
         self._check_adapter_valid(adapter_name)
-        with self.multi_adapter.save_context(kwargs.get('adapter_name')):
+        with self._adapter_context(adapter_name), self.multi_adapter.save_context(adapter_name):
             load_optimizer = kwargs.get('load_optimizer', False)
             if output_dir is None:
                 if os.path.exists(name):

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

@@ -232,9 +232,11 @@ def load_full_state_dict(self, model, state_dict) -> None:
         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
+
+            from .weight_state import with_tied_weight_aliases
             set_model_state_dict(
                 model,
-                state_dict,
+                with_tied_weight_aliases(self.unwrap_model(model), state_dict),
                 options=self._prepare_full_optimizer_state_dict_options(for_load=True),
             )
             return

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

@@ -113,9 +113,10 @@ def _build_ep_fsdp_device_mesh(self, ep_size: Optional[int] = None) -> Optional[
         )
         return ep_mesh.to_torch_device_mesh()
 
-    def wrap_model(self, model, optimizer=None):
+    def wrap_model(self, model, optimizer=None, *additional_optimizers):
+        optimizers = (optimizer, *additional_optimizers)
         if self.device_mesh is None:
-            return model, optimizer
+            return (model, *optimizers)
         fsdp_mesh = _build_fsdp_mesh(self.device_mesh)
         if fsdp_mesh is None:
             # FSDP has nothing to shard with a single rank, but callers still
@@ -126,8 +127,9 @@ def wrap_model(self, model, optimizer=None):
             ep_enabled = (self.enable_ep and self.ep_fsdp_device_mesh is not None)
 
             # Drop optimizer references to pre-shard params before fully_shard to reduce peak memory.
-            if optimizer is not None:
-                _unbind_optimizer_params(optimizer)
+            for optimizer in optimizers:
+                if optimizer is not None:
+                    _unbind_optimizer_params(optimizer)
 
             use_meta = self.use_rank0_pretrained_broadcast()
 
@@ -235,10 +237,8 @@ def wrap_model(self, model, optimizer=None):
             if ep_enabled:
                 _rebuild_ep_param_groups(model)
 
-        if optimizer is not None:
-            optimizer = _rebind_optimizer(optimizer, model)
-
-        return model, optimizer
+        return (model, *(_rebind_optimizer(optimizer, model) if optimizer is not None else None
+                         for optimizer in optimizers))
 
     def _prepare_optimizer_state_dict_options(self, *, for_load: bool):
         from torch.distributed.checkpoint.state_dict import StateDictOptions
@@ -370,9 +370,11 @@ def load_full_state_dict(self, model, state_dict) -> None:
         """
         if self.device_mesh is not None:
             from torch.distributed.checkpoint.state_dict import StateDictOptions, set_model_state_dict
+
+            from .weight_state import with_tied_weight_aliases
             set_model_state_dict(
                 model,
-                state_dict,
+                with_tied_weight_aliases(self.unwrap_model(model), state_dict),
                 options=StateDictOptions(full_state_dict=True, broadcast_from_rank0=True),
             )
             return
@@ -707,7 +709,7 @@ def _rebind_optimizer(optimizer: torch.optim.Optimizer, model: nn.Module) -> tor
                            'or reinitialize it before training.')
     name_to_param = dict(model.named_parameters())
     ep_patched = any(getattr(module, '_ep_patched', False) for module in model.modules())
-    if len(optimizer.param_groups) != 1:
+    if len(optimizer.param_groups) != 1 or 'param_names' in optimizer.param_groups[0]:
         for group in optimizer.param_groups:
             if 'param_names' not in group:
                 raise RuntimeError('NativeFSDPStrategy cannot rebind optimizer param_groups without param_names. '

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

@@ -0,0 +1,17 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+
+
+def with_tied_weight_aliases(model, state_dict: dict) -> dict:
+    """Expand omitted checkpoint names only for parameters shared by identity."""
+    aliases = {}
+    for name, parameter in model.named_parameters(remove_duplicate=False):
+        aliases.setdefault(id(parameter), []).append(name)
+    expanded = dict(state_dict)
+    for names in aliases.values():
+        if len(names) < 2:
+            continue
+        source = next((name for name in names if name in state_dict), None)
+        if source is not None:
+            for name in names:
+                expanded.setdefault(name, state_dict[source])
+    return expanded

src/twinkle/model/transformers/transformers.py

@@ -420,13 +420,12 @@ def _lazy_wrap_model(self):
                 self.sp_strategy.initialize()
             self._ensure_gdn_padding_free_patch()
 
-            if len(optimizer_groups) == 1:
-                optimizer_group = optimizer_groups[0]
-                optimizer = optimizer_group.optimizer
-                assert optimizer is not None
-                self.model, optimizer = self.strategy.wrap_model(self.model, optimizer)
-                optimizer_group.optimizer = optimizer
-                self.register_mm_forward_hook(optimizer_group)
+            if optimizer_groups:
+                result = self.strategy.wrap_model(self.model, *(group.optimizer for group in optimizer_groups))
+                self.model = result[0]
+                for optimizer_group, optimizer in zip(optimizer_groups, result[1:], strict=True):
+                    optimizer_group.optimizer = optimizer
+                    self.register_mm_forward_hook(optimizer_group)
             else:
                 # maybe forward_only, no optimizer_group available
                 result = self.strategy.wrap_model(self.model)
@@ -440,7 +439,24 @@ def register_mm_forward_hook(self, optimizer_group: OptimizerGroup):
         model = self.strategy.unwrap_model(self.model)
         template = optimizer_group.template
         assert template is not None
-        optimizer_group._handler = model.register_forward_pre_hook(template.pre_forward_hook, with_kwargs=True)
+
+        def forward_hook(model, args, kwargs):
+            selected_group = getattr(self, '_forward_optimizer_group', None)
+            if selected_group is None:
+                selected_group = self.optimizer_group.get(self._get_default_group())
+            if selected_group is optimizer_group:
+                return template.pre_forward_hook(model, args, kwargs)
+
+        optimizer_group._handler = model.register_forward_pre_hook(forward_hook, with_kwargs=True)
+
+    @contextlib.contextmanager
+    def _mm_forward_context(self, optimizer_group: OptimizerGroup):
+        previous = getattr(self, '_forward_optimizer_group', None)
+        self._forward_optimizer_group = optimizer_group
+        try:
+            yield
+        finally:
+            self._forward_optimizer_group = previous
 
     def unregister_mm_forward_hook(self, optimizer_group: OptimizerGroup):
         if optimizer_group._handler is not None:
@@ -605,7 +621,7 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec
         rr_cleanup = self._router_replay_setup(router_replay_action, routed_experts, batch_size,
                                                router_replay_manual_cleanup)
 
-        with _resolve_task_context(self.model, task):
+        with _resolve_task_context(self.model, task), self._mm_forward_context(optimizer_config):
             outputs = self.model(**inputs)
 
         recorded_routing = rr_cleanup()
@@ -734,7 +750,7 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T
             lora_ctx = (
                 unwrapped_model.disable_adapter()
                 if disable_lora and isinstance(unwrapped_model, PeftModel) else contextlib.nullcontext())
-            with _resolve_task_context(self.model, task), lora_ctx:
+            with _resolve_task_context(self.model, task), lora_ctx, self._mm_forward_context(optimizer_config):
                 outputs = self.model(**inputs)
 
             recorded_routing = rr_cleanup()
@@ -1598,11 +1614,7 @@ def _restore_training_state(self, checkpoint_dir, *, adapter_name=''):
     def resume_from_checkpoint(self, checkpoint_dir, *, resume_only_model=False, **kwargs):
         adapter_name = kwargs.get('adapter_name', '')
 
-        has_adapter = (
-            os.path.exists(os.path.join(checkpoint_dir, 'adapter_model.safetensors'))
-            or os.path.exists(os.path.join(checkpoint_dir, 'adapter_model.bin')))
-        if has_adapter:
-            self.load(checkpoint_dir, adapter_name=adapter_name)
+        self.load(checkpoint_dir, adapter_name=adapter_name)
 
         if not resume_only_model:
             trainer_state = self._restore_training_state(checkpoint_dir, adapter_name=adapter_name)