From 19d1849ac079c3613ec65cad5374c7a00737ee84 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Wed, 14 Oct 2026 14:32:34 +0800 Subject: [PATCH 1/4] fix: preserve training state across checkpoints and shared adapters --- src/twinkle/model/megatron/megatron.py | 9 +- .../model/megatron/multi_lora_megatron.py | 3 + src/twinkle/model/multi_lora.py | 72 +++------- .../transformers/multi_lora_transformers.py | 7 +- .../model/transformers/strategy/accelerate.py | 4 +- .../transformers/strategy/native_fsdp.py | 22 +-- .../strategy/sequence_parallel/__init__.py | 4 +- .../transformers/strategy/weight_state.py | 17 +++ .../model/transformers/transformers.py | 19 +-- .../model/test_multi_lora_tensor_parallel.py | 129 ++++++++++++++++++ .../test_full_checkpoint_resume.py | 40 ++++++ .../test_megatron_resume_response.py | 49 +++++++ .../test_megatron_weight_only_resume.py | 55 ++++++++ .../test_multi_lora_checkpoint_context.py | 48 +++++++ .../transformers/test_native_fsdp_strategy.py | 34 +++++ .../transformers/test_tied_weight_restore.py | 46 +++++++ 16 files changed, 476 insertions(+), 82 deletions(-) create mode 100644 src/twinkle/model/transformers/strategy/weight_state.py create mode 100644 tests/model/test_multi_lora_tensor_parallel.py create mode 100644 tests/transformers/test_full_checkpoint_resume.py create mode 100644 tests/transformers/test_megatron_resume_response.py create mode 100644 tests/transformers/test_megatron_weight_only_resume.py create mode 100644 tests/transformers/test_multi_lora_checkpoint_context.py create mode 100644 tests/transformers/test_tied_weight_restore.py diff --git a/src/twinkle/model/megatron/megatron.py b/src/twinkle/model/megatron/megatron.py index e48243b15..a782f9270 100644 --- a/src/twinkle/model/megatron/megatron.py +++ b/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()) diff --git a/src/twinkle/model/megatron/multi_lora_megatron.py b/src/twinkle/model/megatron/multi_lora_megatron.py index 8af62df4b..de37d1598 100644 --- a/src/twinkle/model/megatron/multi_lora_megatron.py +++ b/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) diff --git a/src/twinkle/model/multi_lora.py b/src/twinkle/model/multi_lora.py index ca7a5ae79..797c3e628 100644 --- a/src/twinkle/model/multi_lora.py +++ b/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) diff --git a/src/twinkle/model/transformers/multi_lora_transformers.py b/src/twinkle/model/transformers/multi_lora_transformers.py index c8df3f61a..ea889f4ba 100644 --- a/src/twinkle/model/transformers/multi_lora_transformers.py +++ b/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): diff --git a/src/twinkle/model/transformers/strategy/accelerate.py b/src/twinkle/model/transformers/strategy/accelerate.py index 56f878a92..db8983709 100644 --- a/src/twinkle/model/transformers/strategy/accelerate.py +++ b/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 diff --git a/src/twinkle/model/transformers/strategy/native_fsdp.py b/src/twinkle/model/transformers/strategy/native_fsdp.py index a1596bb2e..d3b49fe7f 100644 --- a/src/twinkle/model/transformers/strategy/native_fsdp.py +++ b/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. ' diff --git a/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py b/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py index 2c5bc0c0c..a7c927482 100644 --- a/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py +++ b/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py @@ -1059,9 +1059,9 @@ def gather_loss_tensors( outputs['entropies'] = self._trim_gathered_sequence_padding(gathered_entropies, real_position_ids) return inputs, outputs - def wrap_model(self, model, optimizer=None): + def wrap_model(self, model, optimizer=None, *additional_optimizers): self.initialize() - return model, optimizer + return (model, optimizer, *additional_optimizers) def unwrap_model(self, model): return model diff --git a/src/twinkle/model/transformers/strategy/weight_state.py b/src/twinkle/model/transformers/strategy/weight_state.py new file mode 100644 index 000000000..6ac607d9d --- /dev/null +++ b/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 diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index d209b82a9..2931d5bfb 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/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) @@ -1598,11 +1597,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) diff --git a/tests/model/test_multi_lora_tensor_parallel.py b/tests/model/test_multi_lora_tensor_parallel.py new file mode 100644 index 000000000..fa0867cc9 --- /dev/null +++ b/tests/model/test_multi_lora_tensor_parallel.py @@ -0,0 +1,129 @@ +"""Compare the patched LoRA forward with a dense oracle using real Gloo collectives.""" +import sys +from types import ModuleType, SimpleNamespace + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from torch import nn + + +class CopyToParallel(torch.autograd.Function): + @staticmethod + def forward(ctx, value): + return value + + @staticmethod + def backward(ctx, grad): + grad = grad.clone() + dist.all_reduce(grad) + return grad + + +class ReduceFromParallel(torch.autograd.Function): + @staticmethod + def forward(ctx, value): + value = value.clone() + dist.all_reduce(value) + return value + + @staticmethod + def backward(ctx, grad): + return grad + + +class ParallelLinear(nn.Module): + def __init__(self, weight, mode): + super().__init__() + self.weight = nn.Parameter(weight.clone()) + self.parallel_mode = mode + + def forward(self, inputs, *args, **kwargs): + if self.parallel_mode == 'column': + inputs = CopyToParallel.apply(inputs) + outputs = nn.functional.linear(inputs, self.weight) + if self.parallel_mode == 'row': + outputs = ReduceFromParallel.apply(outputs) + return outputs, None + + +class BaseLinear(ParallelLinear): + def forward(self, inputs, *args, **kwargs): + return inputs.new_zeros(inputs.shape[0], self.weight.shape[0]), None + + +class LoraParallelLinear(nn.Module): + def __init__(self, rank, mode, a, b): + super().__init__() + self.sequence_parallel = self.disable_adapters = self.merged = False + self.config = SimpleNamespace(fp8_param=False) + self.active_adapters = ['slot'] + shard = slice(rank * 2, rank * 2 + 2) + self.base_layer = BaseLinear(torch.zeros(4 if mode == 'row' else 2, + 2 if mode == 'row' else 4), mode) + self.lora_A = nn.ModuleDict({'slot': ParallelLinear(a[:, shard] if mode == 'row' else a, + 'row' if mode == 'row' else None)}) + self.lora_B = nn.ModuleDict({'slot': ParallelLinear(b if mode == 'row' else b[shard], + None if mode == 'row' else 'column')}) + self.lora_dropout = nn.ModuleDict({'slot': nn.Identity()}) + + +def _worker(rank, rendezvous, mode): + torch.set_num_threads(1) + dist.init_process_group('gloo', init_method=rendezvous, rank=rank, world_size=2) + try: + from twinkle.model.multi_lora import MultiLora + + # Only transport-specific TE classes are replaced; the production patch and + # CPU collective/autograd operations below run normally in two processes. + bridge = ModuleType('mcore_bridge') + bridge.LoraParallelLinear = LoraParallelLinear + sys.modules['mcore_bridge'] = bridge + te = ModuleType('megatron.core.extensions.transformer_engine') + te.TELinear = ParallelLinear + te.TEGroupedLinear = type('Grouped', (nn.Module,), {}) + te.TELayerNormColumnParallelLinear = type('LayerNormColumn', (nn.Module,), {}) + sys.modules[te.__name__] = te + tp = ModuleType('megatron.core.tensor_parallel') + tp.gather_from_sequence_parallel_region = tp.scatter_to_sequence_parallel_region = lambda value: value + sys.modules[tp.__name__] = tp + router = ModuleType('megatron.core.transformer.moe.router') + router.TopKRouter = type('Router', (nn.Module,), {}) + sys.modules[router.__name__] = router + torch.manual_seed(42) + a, b = torch.randn(4, 4), torch.randn(4, 4) + reference_a, reference_b = nn.Parameter(a[:2].clone()), nn.Parameter(b[:, :2].clone()) + reference_optim = torch.optim.Adam([reference_a, reference_b], lr=0.01) + layer = LoraParallelLinear(rank, mode, a, b) + config = SimpleNamespace(r=2, lora_alpha=2, use_rslora=False, target_modules='all-linear') + tenant = SimpleNamespace(tenant_config=config, adapter_name='slot') + manager = SimpleNamespace(max_r=4, find_lora=lambda name: tenant, match_target_modules=lambda *args: True) + MultiLora._patch_lora_forward(manager, 'linear', layer) + optimizer = torch.optim.Adam(layer.parameters(), lr=0.01) + inputs, target = torch.randn(2, 4), torch.randn(2, 4) + shard = slice(rank * 2, rank * 2 + 2) + for _ in range(3): + expected = nn.functional.linear(nn.functional.linear(inputs, reference_a), reference_b) + actual, _ = layer(inputs[:, shard] if mode == 'row' else inputs) + torch.testing.assert_close(actual, expected if mode == 'row' else expected[:, shard]) + (expected - target).square().sum().backward() + (actual - (target if mode == 'row' else target[:, shard])).square().sum().backward() + expected_a = reference_a.grad[:, shard] if mode == 'row' else reference_a.grad + expected_b = reference_b.grad if mode == 'row' else reference_b.grad[shard] + torch.testing.assert_close(layer.lora_A['slot'].weight.grad[:2], expected_a) + torch.testing.assert_close(layer.lora_B['slot'].weight.grad[:, :2], expected_b) + assert not torch.count_nonzero(layer.lora_A['slot'].weight.grad[2:]) + assert not torch.count_nonzero(layer.lora_B['slot'].weight.grad[:, 2:]) + reference_optim.step() + optimizer.step() + reference_optim.zero_grad() + optimizer.zero_grad() + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize('mode', ['row', 'column']) +def test_multi_lora_preserves_tensor_parallel_forward_and_gradients(tmp_path, mode): + rendezvous = (tmp_path / 'gloo').as_uri() + mp.spawn(_worker, args=(rendezvous, mode), nprocs=2, join=True) diff --git a/tests/transformers/test_full_checkpoint_resume.py b/tests/transformers/test_full_checkpoint_resume.py new file mode 100644 index 000000000..63ae4b99e --- /dev/null +++ b/tests/transformers/test_full_checkpoint_resume.py @@ -0,0 +1,40 @@ +import json +from types import SimpleNamespace + +import pytest +import torch +from safetensors.torch import save_file + +from twinkle.model.transformers.strategy.native_fsdp import NativeFSDPStrategy +from twinkle.model.transformers.transformers import TransformersModel + + +@pytest.mark.parametrize('weights_only', [True, False]) +def test_full_resume_restores_weights_and_optional_training_state(tmp_path, weights_only): + model = torch.nn.Linear(1, 1, bias=False) + model.weight.data.fill_(7) + optimizer = torch.optim.Adam(model.parameters(), lr=0.1) + source = torch.nn.Parameter(torch.ones(1, 1)) + saved_optimizer = torch.optim.Adam([source], lr=0.1) + source.grad = torch.full_like(source, 5) + saved_optimizer.step() + torch.save(saved_optimizer.state_dict(), tmp_path / 'optimizer.pt') + save_file({'weight': torch.full((1, 1), 2.)}, str(tmp_path / 'model.safetensors')) + (tmp_path / 'trainer_state.json').write_text(json.dumps({ + 'cur_step': 3, 'consumed_train_samples': 12, 'gradient_accumulation_steps': 2, + })) + wrapper = object.__new__(TransformersModel) + wrapper.__dict__['model'] = model + wrapper.strategy = NativeFSDPStrategy(device_mesh=None, enable_ep=False) + wrapper.optimizer_group = {'': SimpleNamespace(optimizer=optimizer, lr_scheduler=None, scaler=None, + cur_step=0, gradient_accumulation_steps=1)} + result = wrapper.resume_from_checkpoint(str(tmp_path), resume_only_model=weights_only) + assert model(torch.ones(1, 1)).item() == 2 + assert result['cur_step'] == 3 + if weights_only: + assert not optimizer.state + assert wrapper.optimizer_group[''].cur_step == 0 + else: + torch.testing.assert_close(optimizer.state[model.weight]['exp_avg'], torch.full((1, 1), 0.5)) + assert wrapper.optimizer_group[''].cur_step == 3 + assert wrapper.optimizer_group[''].gradient_accumulation_steps == 2 diff --git a/tests/transformers/test_megatron_resume_response.py b/tests/transformers/test_megatron_resume_response.py new file mode 100644 index 000000000..f5e679cd6 --- /dev/null +++ b/tests/transformers/test_megatron_resume_response.py @@ -0,0 +1,49 @@ +from types import SimpleNamespace + +import pytest +import ray + +import twinkle.infra as infra +from twinkle.model.megatron.megatron import MegatronModel +from twinkle.model.megatron.multi_lora_megatron import MultiLoraMegatronModel +from twinkle.protocol.types.model import TrainingProgressResponse + + +@ray.remote +class ProgressWorker: + def __init__(self): + self.calls = 0 + + def resume_from_checkpoint(self, path, **kwargs): + self.calls += 1 + return {'cur_step': 4, 'consumed_train_samples': 16, 'gradient_accumulation_steps': 1} + + def call_count(self): + return self.calls + + +@pytest.fixture(scope='module') +def workers(): + owns_runtime = not ray.is_initialized() + if owns_runtime: + ray.init(address='local', num_cpus=2, include_dashboard=False, log_to_driver=False) + actors = [ProgressWorker.remote() for _ in range(2)] + try: + yield actors + finally: + for actor in actors: + ray.kill(actor, no_restart=True) + if owns_runtime: + ray.shutdown() + + +@pytest.mark.parametrize('model_class', [MegatronModel, MultiLoraMegatronModel]) +def test_megatron_resume_returns_one_typed_progress_after_calling_all_workers(workers, monkeypatch, model_class): + monkeypatch.setattr(infra, '_mode', 'ray') + monkeypatch.setattr(infra, '_lazy_collect', False) + before = ray.get([worker.call_count.remote() for worker in workers]) + driver = SimpleNamespace(_actors=workers, _max_concurrency=None, device_mesh=None) + result = model_class.resume_from_checkpoint(driver, 'checkpoint') + progress = TrainingProgressResponse.model_validate({'result': result}) + assert progress.result['cur_step'] == 4 + assert ray.get([worker.call_count.remote() for worker in workers]) == [count + 1 for count in before] diff --git a/tests/transformers/test_megatron_weight_only_resume.py b/tests/transformers/test_megatron_weight_only_resume.py new file mode 100644 index 000000000..1fda7208d --- /dev/null +++ b/tests/transformers/test_megatron_weight_only_resume.py @@ -0,0 +1,55 @@ +from contextlib import nullcontext +from copy import deepcopy +from types import SimpleNamespace + +import pytest +import torch + +from twinkle.model.megatron.megatron import MegatronModel +from twinkle.model.megatron.multi_lora_megatron import MultiLoraMegatronModel + + +class MasterParameterOptimizer: + """Use real Adam with the documented MCore model/master parameter contract.""" + def __init__(self, model_parameter): + self.model_parameter = model_parameter + self.master = torch.nn.Parameter(model_parameter.detach().float().clone()) + self.adam = torch.optim.Adam([self.master], lr=0.1) + + def reload_model_params(self): + self.master.data.copy_(self.model_parameter) + + def step(self): + self.master.grad = torch.ones_like(self.master) + self.adam.step() + self.model_parameter.data.copy_(self.master.to(self.model_parameter.dtype)) + + +@pytest.mark.parametrize('model_class', [MegatronModel, MultiLoraMegatronModel]) +def test_weights_only_load_refreshes_target_master_params_and_preserves_adam_state(tmp_path, model_class): + parameter = torch.nn.Parameter(torch.tensor([7.], dtype=torch.bfloat16)) + optimizer = MasterParameterOptimizer(parameter) + optimizer.step() + before = deepcopy(optimizer.adam.state_dict()) + expected = MasterParameterOptimizer(torch.nn.Parameter(torch.tensor([2.], dtype=torch.bfloat16))) + expected.adam.load_state_dict(deepcopy(before)) + expected.step() + other = MasterParameterOptimizer(torch.nn.Parameter(torch.tensor([11.], dtype=torch.bfloat16))) + group = SimpleNamespace(optimizer=optimizer) + + def load_weights(*args, **kwargs): + parameter.data.fill_(2.) + + wrapper = SimpleNamespace( + _get_default_group=lambda: 'target', model=[parameter], + strategy=SimpleNamespace(unwrap_model=lambda model: model, + bridge=SimpleNamespace(load_weights=load_weights)), + optimizer_group={'target': group, 'other': SimpleNamespace(optimizer=other)}, + multi_adapter=SimpleNamespace(save_context=lambda adapter: nullcontext('slot'), load_lora_converter=None), + ) + model_class.load(wrapper, str(tmp_path), adapter_name='target', load_optimizer=False) + assert parameter.item() == 2. + assert optimizer.adam.state_dict()['state'][0]['step'] == before['state'][0]['step'] + optimizer.step() + torch.testing.assert_close(parameter, expected.model_parameter) + assert other.master.item() == other.model_parameter.item() == 11. diff --git a/tests/transformers/test_multi_lora_checkpoint_context.py b/tests/transformers/test_multi_lora_checkpoint_context.py new file mode 100644 index 000000000..1b89e3d51 --- /dev/null +++ b/tests/transformers/test_multi_lora_checkpoint_context.py @@ -0,0 +1,48 @@ +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch +from peft import LoraConfig +from torch.distributed.checkpoint.state_dict import get_optimizer_state_dict + +from twinkle.model.multi_lora import MultiLora +from twinkle.model.transformers.multi_lora_transformers import MultiLoraTransformersModel +from twinkle.model.transformers.transformers import TransformersModel + + +@pytest.mark.parametrize('operation', ['save', 'load']) +def test_checkpoint_activates_target_tenant_before_optimizer_state_initialization(operation, tmp_path): + manager = MultiLora(max_loras=2, max_r=4) + config = LoraConfig(r=4, lora_alpha=8, target_modules=['0']) + model = manager.patch(torch.nn.Sequential(torch.nn.Linear(4, 4)), target_modules=['0'], lora_config=config) + manager.acquire_lora('first', config) + manager.acquire_lora('second', config) + wrapper = object.__new__(MultiLoraTransformersModel) + wrapper.__dict__['model'] = model + wrapper.multi_adapter = manager + wrapper.strategy = SimpleNamespace(unwrap_model=lambda value: value) + with manager.adapter('first'): + params = [param for name, param in model.named_parameters() if '.lora_0.' in name] + optimizer = torch.optim.Adam(params) + wrapper.optimizer_group = {'first': SimpleNamespace(optimizer=optimizer), 'second': SimpleNamespace()} + manager.activate_adapter('second') + assert all(not param.requires_grad for param in params) + + def checkpoint_optimizer(*args, **kwargs): + state = get_optimizer_state_dict(model, optimizer) + assert len(state['state']) == len(params) + assert all('exp_avg' in value for value in state['state'].values()) + assert all(param.requires_grad for param in params) + return str(tmp_path) + + if operation == 'save': + with patch.object(TransformersModel, 'save', side_effect=checkpoint_optimizer): + wrapper.save('checkpoint', adapter_name='first', save_optimizer=True) + else: + object.__setattr__(wrapper, '_restore_training_state', checkpoint_optimizer) + with patch('twinkle.model.transformers.multi_lora_transformers.load_peft_weights', return_value={}), \ + patch.object(manager, 'set_state_dict'): + wrapper.load(str(tmp_path), adapter_name='first', load_optimizer=True) + assert model.active_adapter == 'lora_0' + assert set(model.peft_config) == {'lora_0', 'lora_1'} diff --git a/tests/transformers/test_native_fsdp_strategy.py b/tests/transformers/test_native_fsdp_strategy.py index 79402c785..13cb578d9 100644 --- a/tests/transformers/test_native_fsdp_strategy.py +++ b/tests/transformers/test_native_fsdp_strategy.py @@ -1,9 +1,12 @@ # Copyright (c) ModelScope Contributors. All rights reserved. from unittest.mock import MagicMock, patch +from types import SimpleNamespace +import pytest import torch from twinkle.model.transformers.strategy.native_fsdp import NativeFSDPStrategy +from twinkle.model.transformers.transformers import TransformersModel def test_single_rank_native_fsdp_places_model_on_local_device(): @@ -23,3 +26,34 @@ def test_single_rank_native_fsdp_places_model_on_local_device(): model.to.assert_called_once_with(torch.device('cuda:0')) assert wrapped_model is moved_model assert wrapped_optimizer is None + + +@pytest.mark.parametrize('tenant_count', [1, 2]) +def test_lazy_wrapping_rebinds_each_tenants_optimizer_without_other_parameters(tenant_count): + original = torch.nn.ParameterDict({name: torch.nn.Parameter(torch.ones(2)) for name in ('a', 'b', 'base')}) + replacement = torch.nn.ParameterDict({name: torch.nn.Parameter(torch.ones(2)) for name in original}) + groups = { + name: SimpleNamespace(optimizer=torch.optim.Adam([{'params': [original[name]], 'param_names': [name]}], lr=0.1)) + for name in ('a', 'b')[:tenant_count] + } + mesh = MagicMock() + mesh.mesh_dim_names = None + strategy = NativeFSDPStrategy(device_mesh=mesh, enable_ep=False) + wrapper = SimpleNamespace(model=original, optimizer_group=groups, strategy=strategy, + _model_wrapped=False, _enable_expert_parallel=False, sp_strategy=None, + _maybe_apply_expert_parallel=lambda: None, _ensure_sp_strategy=lambda: None, + _ensure_gdn_padding_free_patch=lambda: None, + register_mm_forward_hook=lambda group: None) + with patch.object(original, 'to', return_value=replacement), patch( + 'twinkle.model.transformers.strategy.native_fsdp.Platform.get_local_device', return_value='cpu'): + TransformersModel._lazy_wrap_model(wrapper) + assert wrapper.model is replacement + for name, group in groups.items(): + params = group.optimizer.param_groups[0]['params'] + assert len(params) == 1 and params[0] is replacement[name] + replacement[name].grad = torch.ones_like(replacement[name]) + group.optimizer.step() + assert not torch.equal(replacement[name], original[name]) + torch.testing.assert_close(replacement['base'], torch.ones(2)) + if tenant_count == 1: + torch.testing.assert_close(replacement['b'], torch.ones(2)) diff --git a/tests/transformers/test_tied_weight_restore.py b/tests/transformers/test_tied_weight_restore.py new file mode 100644 index 000000000..e0006cb46 --- /dev/null +++ b/tests/transformers/test_tied_weight_restore.py @@ -0,0 +1,46 @@ +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch + +from twinkle.model.transformers.strategy.accelerate import AccelerateStrategy +from twinkle.model.transformers.strategy.native_fsdp import NativeFSDPStrategy + + +@pytest.mark.parametrize('backend', ['native', 'accelerate']) +@pytest.mark.parametrize('missing_independent_weight', [False, True]) +def test_full_restore_expands_tied_aliases_and_rejects_missing_independent_weights(backend, missing_independent_weight): + model = torch.nn.Module() + model.embedding = torch.nn.Embedding(4, 3) + model.lm_head = torch.nn.Linear(3, 4, bias=False) + model.lm_head.weight = model.embedding.weight + model.other = torch.nn.Linear(3, 3) + saved = {name: value.clone() for name, value in model.state_dict().items()} + checkpoint = {name: value for name, value in saved.items() if name != 'lm_head.weight'} + if missing_independent_weight: + checkpoint.pop('other.weight') + for param in model.parameters(): + param.data.fill_(7) + if backend == 'native': + strategy = object.__new__(NativeFSDPStrategy) + strategy.device_mesh = object() + else: + strategy = object.__new__(AccelerateStrategy) + strategy.accelerator = SimpleNamespace( + unwrap_model=lambda value, **kwargs: value, + state=SimpleNamespace(fsdp_plugin=SimpleNamespace(fsdp_version=2))) + + def strict_cpu_load(target, state, *, options): + return target.load_state_dict(state, strict=options.strict) + + # Replace only distributed transport; the real strict PyTorch loader checks keys and writes weights. + with patch('torch.distributed.checkpoint.state_dict.set_model_state_dict', side_effect=strict_cpu_load): + if missing_independent_weight: + with pytest.raises(RuntimeError, match='other.weight'): + strategy.load_full_state_dict(model, checkpoint) + else: + strategy.load_full_state_dict(model, checkpoint) + for name, actual in model.state_dict().items(): + torch.testing.assert_close(actual, saved[name]) + assert 'lm_head.weight' not in checkpoint From 546d8d1428ab445665bda245c1816b114ae0b5f1 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Wed, 14 Oct 2026 14:52:02 +0800 Subject: [PATCH 2/4] fix: use supported Transformers config imports in Megatron strategy --- src/twinkle/model/megatron/strategy/megatron.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/twinkle/model/megatron/strategy/megatron.py b/src/twinkle/model/megatron/strategy/megatron.py index 59f545a2c..7b8c822bd 100644 --- a/src/twinkle/model/megatron/strategy/megatron.py +++ b/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, ): From 4dc1346638fe523096999f77cc0364e1d636eb79 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Thu, 15 Oct 2026 13:32:19 +0800 Subject: [PATCH 3/4] fix: select one multimodal template hook per forward --- .../model/transformers/transformers.py | 23 +++- .../test_multimodal_group_hooks.py | 108 ++++++++++++++++++ 2 files changed, 128 insertions(+), 3 deletions(-) create mode 100644 tests/transformers/test_multimodal_group_hooks.py diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index 2931d5bfb..b500f65d0 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -439,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: @@ -604,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() @@ -733,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() diff --git a/tests/transformers/test_multimodal_group_hooks.py b/tests/transformers/test_multimodal_group_hooks.py new file mode 100644 index 000000000..701a3c0f3 --- /dev/null +++ b/tests/transformers/test_multimodal_group_hooks.py @@ -0,0 +1,108 @@ +"""Exercise real Qwen3.5 template hooks with multiple optimizer-backed groups.""" +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +from transformers import ProcessorMixin + +from twinkle.model.transformers.transformers import TransformersModel +from twinkle.processor import InputProcessor +from twinkle.template.qwen3_5_vl import Qwen3_5Template + + +class Visual(torch.nn.Module): + dtype = torch.float32 + + def forward(self, values, grid_thw): + return values + + +class VisionLanguageModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.model = torch.nn.Module() + self.model.embed_tokens = torch.nn.Embedding(4, 3) + self.model.visual = Visual() + self.lm_head = torch.nn.Linear(3, 4) + self.config = SimpleNamespace(image_token_id=0, video_token_id=3) + self.fail = False + + def forward(self, inputs_embeds, position_ids=None): + if self.fail: + raise RuntimeError('model forward failed') + return {'logits': self.lm_head(inputs_embeds)} + + +class EncodedProcessor(InputProcessor): + def __call__(self, inputs, **kwargs): + return dict(inputs) + + +@pytest.fixture +def wrapper(): + model = VisionLanguageModel() + wrapper = object.__new__(TransformersModel) + torch.nn.Module.__init__(wrapper) + wrapper.model = model + wrapper.active_group = 'b' + wrapper.device_mesh = None + wrapper.hf_config = model.config + wrapper.sp_strategy = None + wrapper._model_wrapped = False + wrapper._enable_expert_parallel = False + wrapper._gdn_padding_free_enabled = False + wrapper._maybe_apply_expert_parallel = lambda: None + wrapper._ensure_sp_strategy = lambda: None + wrapper._ensure_gdn_padding_free_patch = lambda: None + wrapper.strategy = SimpleNamespace( + capture_pre_ep_state_if_needed=lambda *args, **kwargs: None, + wrap_model=lambda model, *optimizers: (model, *optimizers), + unwrap_model=lambda model: model, + ) + groups = {} + for name, parameter in zip(('a', 'b'), model.lm_head.parameters()): + template = object.__new__(Qwen3_5Template) + template.processor = object.__new__(ProcessorMixin) + template.pre_forward_hook = Mock(wraps=template.pre_forward_hook) + groups[name] = SimpleNamespace( + optimizer=torch.optim.Adam([parameter]), template=template, processor=EncodedProcessor(), + loss_instance=SimpleNamespace(require_logps=False), train_status=SimpleNamespace(), + eval_status=SimpleNamespace(), accumulate_metrics=lambda training: None, _handler=None, + ) + wrapper.optimizer_group = groups + return wrapper + + +def inputs(): + return {'input_ids': torch.tensor([[0, 1]]), 'pixel_values': torch.tensor([[1., 2., 3.]]), + 'image_grid_thw': torch.tensor([[1, 1, 1]])} + + +@pytest.mark.parametrize('method', ['forward', 'forward_only']) +@pytest.mark.parametrize('first', ['a', 'b']) +def test_only_requested_multimodal_group_processes_each_forward(wrapper, method, first): + expected_calls = {'a': 0, 'b': 0} + for name in (first, 'b' if first == 'a' else 'a', first): + result = getattr(wrapper, method)(inputs=inputs(), adapter_name=name, return_logits=True) + assert result['logits'].shape == (1, 2, 4) + assert result['logits'].requires_grad == (method == 'forward') + expected_calls[name] += 1 + for identity, group in wrapper.optimizer_group.items(): + assert group.template.pre_forward_hook.call_count == expected_calls[identity] + assert wrapper.active_group == 'b' + + +@pytest.mark.parametrize('method', ['forward', 'forward_only']) +def test_failed_forward_does_not_leave_another_groups_template_selected(wrapper, method): + wrapper.model.fail = True + with pytest.raises(RuntimeError, match='model forward failed'): + getattr(wrapper, method)(inputs=inputs(), adapter_name='a') + wrapper.model.fail = False + assert wrapper.model(**inputs())['logits'].shape == (1, 2, 4) + assert wrapper.optimizer_group['a'].template.pre_forward_hook.call_count == 1 + assert wrapper.optimizer_group['b'].template.pre_forward_hook.call_count == 1 + result = getattr(wrapper, method)(inputs=inputs(), return_logits=True) + assert result['logits'].shape == (1, 2, 4) + assert wrapper.optimizer_group['a'].template.pre_forward_hook.call_count == 1 + assert wrapper.optimizer_group['b'].template.pre_forward_hook.call_count == 2 From 9e855f03d04d04770cad18985fcbe48b07dcdaeb Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Thu, 15 Oct 2026 14:12:56 +0800 Subject: [PATCH 4/4] fix: drop unused sequence parallel optimizer interface expansion --- .../model/transformers/strategy/sequence_parallel/__init__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py b/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py index a7c927482..2c5bc0c0c 100644 --- a/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py +++ b/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py @@ -1059,9 +1059,9 @@ def gather_loss_tensors( outputs['entropies'] = self._trim_gathered_sequence_padding(gathered_entropies, real_position_ids) return inputs, outputs - def wrap_model(self, model, optimizer=None, *additional_optimizers): + def wrap_model(self, model, optimizer=None): self.initialize() - return (model, optimizer, *additional_optimizers) + return model, optimizer def unwrap_model(self, model): return model