From 719313b22bad86e55038038d801f1685bbdedc1c Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Wed, 12 Aug 2026 17:12:54 +0800 Subject: [PATCH 01/28] fix(transformers): support memory-efficient multi-LoRA EP init --- .../ep_fsdp2_multi_lora_deepseek_v4.py | 26 +++-- src/twinkle/model/multi_lora.py | 26 ++++- .../model/multi_lora_target_parameters.py | 104 ++++++++++++++---- .../model/transformers/moe/expert_parallel.py | 12 +- .../transformers/multi_lora_transformers.py | 29 +++-- .../transformers/strategy/native_fsdp.py | 15 ++- .../test_ep_multi_lora_target_parameters.py | 98 ++++++++++++++++- 7 files changed, 254 insertions(+), 56 deletions(-) diff --git a/cookbook/transformers/ep_fsdp2_multi_lora_deepseek_v4.py b/cookbook/transformers/ep_fsdp2_multi_lora_deepseek_v4.py index bc06ef5f2..799cd1111 100644 --- a/cookbook/transformers/ep_fsdp2_multi_lora_deepseek_v4.py +++ b/cookbook/transformers/ep_fsdp2_multi_lora_deepseek_v4.py @@ -1,8 +1,8 @@ # Copyright (c) ModelScope Contributors. All rights reserved. """EP + FSDP2 + Multi-LoRA SFT cookbook for DeepSeek-V4. -Run on 8 GPUs: - torchrun --nproc-per-node=8 cookbook/transformers/ep_fsdp2_multi_lora_deepseek_v4.py +Run on 2 GPUs: + torchrun --nproc-per-node=2 cookbook/transformers/ep_fsdp2_multi_lora_deepseek_v4.py """ import os from pathlib import Path @@ -19,10 +19,10 @@ logger = get_logger() -MODEL_ID = os.environ.get('DSV4_MODEL_ID', 'ms://deepseek-ai/DeepSeek-V4-Flash') -DATASET_ID = os.environ.get('DATASET_ID', 'ms://swift/self-cognition') +MODEL_ID = os.environ.get('DSV4_MODEL_ID', '/nas/disk1/random-deepseek-v4-4b') +DATASET_ID = os.environ.get('DATASET_ID', '/model/ljl/dataset/self-cognition.jsonl') TEMPLATE_ID = os.environ.get('TEMPLATE_ID', 'DeepseekV4Template') -BATCH_SIZE = int(os.environ.get('BATCH_SIZE', '4')) +BATCH_SIZE = int(os.environ.get('BATCH_SIZE', '1')) GRAD_ACCUM_STEPS = int(os.environ.get('GRAD_ACCUM_STEPS', '4')) LOG_INTERVAL = GRAD_ACCUM_STEPS LR = float(os.environ.get('LR', '1e-4')) @@ -37,11 +37,12 @@ RESUME_ONLY_MODEL = os.environ.get('RESUME_ONLY_MODEL', '0') == '1' IGNORE_DATA_SKIP = os.environ.get('IGNORE_DATA_SKIP', '0') == '1' ADAPTER_NAMES = [name.strip() for name in os.environ.get('ADAPTER_NAMES', 'tenant_a,tenant_b').split(',') if name] +WORLD_SIZE = int(os.environ.get('WORLD_SIZE', '2')) device_mesh = DeviceMesh.from_sizes( - fsdp_size=8, + fsdp_size=WORLD_SIZE, dp_size=1, - ep_size=8, + ep_size=WORLD_SIZE, device_type=Platform.get_platform().device_prefix(), ) twinkle.initialize(mode='local', global_device_mesh=device_mesh) @@ -86,14 +87,14 @@ def train(): dataset.encode(batched=True) dataloader = DataLoader(dataset=dataset, batch_size=BATCH_SIZE, device_mesh=device_mesh) - ep_lora_cfg = _build_lora_config(enable_ep=ENABLE_EP) # LoraConfig for target params + ep_lora_cfg = _build_lora_config(enable_ep=ENABLE_EP) # LoraConfig for target params lora_cfg = _build_lora_config(enable_ep=False) # LoraConfig for PEFT adapter model = MultiLoraTransformersModel( model_id=MODEL_ID, config=config, device_mesh=device_mesh, strategy='native_fsdp', - memory_efficient_init=False, + memory_efficient_init=True, max_loras=MAX_LORAS, max_r=MAX_R, fsdp_config={ @@ -109,6 +110,10 @@ def train(): for adapter_name in ADAPTER_NAMES: model.add_adapter_to_model(adapter_name, ep_lora_cfg, gradient_accumulation_steps=GRAD_ACCUM_STEPS) + # Materialize rank-local EP/FSDP shards before loading adapter checkpoints + # or creating optimizers. + model._lazy_wrap_model() + if RESUME_FROM_CHECKPOINT: checkpoint_path = Path(RESUME_FROM_CHECKPOINT).expanduser().resolve() progress = None @@ -128,9 +133,6 @@ def train(): f'Total steps: {len(dataloader)}, batch_size={BATCH_SIZE}, grad_accum={GRAD_ACCUM_STEPS}, ' f'enable_ep={ENABLE_EP}, adapters={ADAPTER_NAMES}, output_dir={OUTPUT_DIR}') - # After LoRA init, before forward (LoRA active): perform EP + FSDP broadcast & sharding. - model._lazy_wrap_model() - # Must call set_optimizer() after EP + FSDP sharding, otherwise optimizer may # capture stale parameter references and fail to update the actual LoRA weights. for adapter_name in ADAPTER_NAMES: diff --git a/src/twinkle/model/multi_lora.py b/src/twinkle/model/multi_lora.py index f310306c6..0a66b81ca 100644 --- a/src/twinkle/model/multi_lora.py +++ b/src/twinkle/model/multi_lora.py @@ -30,14 +30,18 @@ class LoraTenant: class MultiLora: - def __init__(self, max_loras=5, max_r=32, max_length: int = 8192): + def __init__(self, max_loras=5, max_r=32, max_length: int = 8192, defer_initial_weights: bool = False): self.max_loras = max_loras self.max_r = max_r self.loras: List[LoraTenant] = [] self.module: PeftModel self._active_adapters = [] self.max_length = max_length - self.target_parameter_manager = TargetParameterLoraManager(max_loras=max_loras, max_r=max_r) + self.target_parameter_manager = TargetParameterLoraManager( + max_loras=max_loras, + max_r=max_r, + defer_initial_weights=defer_initial_weights, + ) def _get_available_lora(self) -> Optional[LoraTenant]: for _lora in self.loras: @@ -48,6 +52,12 @@ def _get_available_lora(self) -> Optional[LoraTenant]: def _read_param_tensor(self, parameter): return torch_util.to_local_tensor(parameter) + @staticmethod + def _read_local_param_tensor(parameter): + if hasattr(parameter, 'to_local'): + return parameter.to_local() + return parameter + @staticmethod def _is_distributed_param(parameter): return hasattr(parameter, 'device_mesh') and hasattr(parameter, 'placements') @@ -568,13 +578,19 @@ def _store_weights(_module): if self._is_target_parameter_lora_name(name): continue if pattern.search(name): - lora_tenant.lora_A_weights[name] = self._read_param_tensor(parameter).clone().to('cpu') + local_parameter = self._read_local_param_tensor(parameter) + if local_parameter.is_meta: + raise RuntimeError( + f'LoRA parameter {name} is still on meta; materialize the model before saving ' + 'its initial weights.') + lora_tenant.lora_A_weights[name] = local_parameter.detach().cpu().clone() if isinstance(self.module, list): for _module in self.module: _store_weights(_module) else: _store_weights(self.module) + self.target_parameter_manager.save_initial_weights() def load_lora_converter(self, name, parameter, **kwargs): @@ -741,12 +757,12 @@ def _load_initial_weights(_module): if self._is_target_parameter_lora_name(name): continue if pattern_A.search(name): - local_param = self._read_param_tensor(parameter) + local_param = self._read_local_param_tensor(parameter) if local_param is not None: value = _lora.lora_A_weights[name].to(dtype=parameter.dtype, device=local_param.device) self._write_param_tensor(parameter, value) if pattern_B.search(name): - local_param = self._read_param_tensor(parameter) + local_param = self._read_local_param_tensor(parameter) if local_param is not None: self._write_param_tensor(parameter, torch.zeros_like(local_param)) diff --git a/src/twinkle/model/multi_lora_target_parameters.py b/src/twinkle/model/multi_lora_target_parameters.py index 8955917da..d6936776c 100644 --- a/src/twinkle/model/multi_lora_target_parameters.py +++ b/src/twinkle/model/multi_lora_target_parameters.py @@ -37,7 +37,13 @@ def key(self) -> str: class TargetParameterLoraWrapper(nn.Module): - def __init__(self, record: TargetParameterRecord, max_loras: int, max_r: int): + def __init__( + self, + record: TargetParameterRecord, + max_loras: int, + max_r: int, + defer_initial_weights: bool = False, + ): super().__init__() self.record = record # Unsharded original target parameter (pre-sharding snapshot) @@ -53,6 +59,7 @@ def __init__(self, record: TargetParameterRecord, max_loras: int, max_r: int): self.max_loras = max_loras self.max_r = max_r + self.defer_initial_weights = defer_initial_weights self.active_adapter: str | None = None self.disable_adapters = False self.lora_A = nn.ParameterDict() @@ -90,10 +97,7 @@ def _init_slots(self) -> None: raise ValueError( f'target parameter {self.record.key} has {parameter.ndim} dimensions; only 2D and 3D are supported') - # Note: reset_slot requires the tensor to be created on a physical device, not on a meta device. device = parameter.device - if device.type == 'meta': - device = 'cpu' for index in range(self.max_loras): slot_name = f'lora_{index}' self.lora_A[slot_name] = nn.Parameter( @@ -114,24 +118,70 @@ def _init_slots(self) -> None: )) self.r[slot_name] = self.max_r self.scaling[slot_name] = 1.0 - self.reset_slot(slot_name) + if device.type != 'meta': + nn.init.kaiming_uniform_(self.lora_A[slot_name], a=math.sqrt(5)) + nn.init.zeros_(self.lora_B[slot_name]) + if not self.defer_initial_weights: + self._initial_lora_A[slot_name] = self.lora_A[slot_name].detach().cpu().clone() + + @staticmethod + def _read_parameter(parameter: nn.Parameter) -> torch.Tensor: + if hasattr(parameter, 'to_local'): + return parameter.to_local() + if hasattr(parameter, 'full_tensor'): + return parameter.full_tensor() + return parameter + + @staticmethod + def _write_parameter(parameter: nn.Parameter, value: torch.Tensor) -> None: + if hasattr(parameter, 'to_local') and hasattr(parameter, 'device_mesh'): + local_parameter = parameter.to_local() + if tuple(value.shape) == tuple(local_parameter.shape): + local_parameter.copy_(value.to(device=local_parameter.device, dtype=local_parameter.dtype)) + return + if tuple(value.shape) != tuple(parameter.shape): + raise ValueError( + f'Cannot restore target-parameter LoRA with shape {tuple(value.shape)} to distributed ' + f'parameter with global shape {tuple(parameter.shape)} and local shape ' + f'{tuple(local_parameter.shape)}') + from torch.distributed.tensor import distribute_tensor + distributed = distribute_tensor( + value.to(device=parameter.device, dtype=parameter.dtype), + parameter.device_mesh, + parameter.placements, + ) + local_parameter.copy_(distributed.to_local()) + return + parameter.copy_(value.to(device=parameter.device, dtype=parameter.dtype)) + + def save_initial_weights(self) -> None: + for slot_name, parameter in self.lora_A.items(): + tensor = self._read_parameter(parameter) + if tensor.is_meta: + raise RuntimeError( + f'Target-parameter LoRA slot {self.record.key}.{slot_name} is still on meta; ' + 'materialize the model before saving its initial weights.') + self._initial_lora_A[slot_name] = tensor.detach().cpu().clone() def reset_slot(self, slot_name: str) -> None: - if slot_name not in self._initial_lora_A: - nn.init.kaiming_uniform_(self.lora_A[slot_name], a=math.sqrt(5)) - self._initial_lora_A[slot_name] = self.lora_A[slot_name].detach().clone().cpu() - else: - initial = self._initial_lora_A[slot_name] - if hasattr(self.record.module, '_ep_local_start') and hasattr(self.record.module, '_ep_local_end'): - start = self.record.module._ep_local_start - end = self.record.module._ep_local_end - initial = initial[start:end] - initial = initial.to( - device=self.lora_A[slot_name].device, - dtype=self.lora_A[slot_name].dtype, - ) - self.lora_A[slot_name].data.copy_(initial) - nn.init.zeros_(self.lora_B[slot_name]) + if self.lora_A[slot_name].is_meta or self.lora_B[slot_name].is_meta: + raise RuntimeError( + f'Target-parameter LoRA slot {self.record.key}.{slot_name} is still on meta; ' + 'materialize the model before resetting it.') + with torch.no_grad(): + if slot_name not in self._initial_lora_A: + nn.init.kaiming_uniform_(self.lora_A[slot_name], a=math.sqrt(5)) + self._initial_lora_A[slot_name] = self.lora_A[slot_name].detach().clone().cpu() + else: + initial = self._initial_lora_A[slot_name] + if (initial.shape[0] != self.lora_A[slot_name].shape[0] + and hasattr(self.record.module, '_ep_local_start') + and hasattr(self.record.module, '_ep_local_end')): + start = self.record.module._ep_local_start + end = self.record.module._ep_local_end + initial = initial[start:end] + self._write_parameter(self.lora_A[slot_name], initial) + self._read_parameter(self.lora_B[slot_name]).zero_() def configure_slot(self, slot_name: str, config: LoraConfig) -> None: if slot_name not in self.lora_A: @@ -280,9 +330,10 @@ def set_state_dict(self, slot_name: str, state_dict: dict[str, torch.Tensor]) -> class TargetParameterLoraManager: - def __init__(self, max_loras: int, max_r: int): + def __init__(self, max_loras: int, max_r: int, defer_initial_weights: bool = False): self.max_loras = max_loras self.max_r = max_r + self.defer_initial_weights = defer_initial_weights self.wrappers: list[TargetParameterLoraWrapper] = [] self.tenant_to_slot: dict[str, str] = {} self.tenant_configs: dict[str, LoraConfig] = {} @@ -312,7 +363,12 @@ def patch(self, model: nn.Module, target_parameters: Iterable[str]) -> None: raise ValueError(f'target_parameters={target_parameters} were set but no parameter was matched') for record in records: - wrapper = TargetParameterLoraWrapper(record, max_loras=self.max_loras, max_r=self.max_r) + wrapper = TargetParameterLoraWrapper( + record, + max_loras=self.max_loras, + max_r=self.max_r, + defer_initial_weights=self.defer_initial_weights, + ) record.module.add_module(f'_twinkle_lora_{record.parameter_name}', wrapper) self.wrappers.append(wrapper) self._assign_peft_key_prefixes() @@ -348,6 +404,10 @@ def release(self, tenant_adapter_name: str) -> None: for wrapper in self.wrappers: wrapper.reset_slot(slot_name) + def save_initial_weights(self) -> None: + for wrapper in self.wrappers: + wrapper.save_initial_weights() + @contextmanager def adapter(self, tenant_adapter_name: str, disable_lora: bool = False): slot_name = self.tenant_to_slot.get(tenant_adapter_name) diff --git a/src/twinkle/model/transformers/moe/expert_parallel.py b/src/twinkle/model/transformers/moe/expert_parallel.py index 218e7b337..aa828d24b 100644 --- a/src/twinkle/model/transformers/moe/expert_parallel.py +++ b/src/twinkle/model/transformers/moe/expert_parallel.py @@ -431,8 +431,10 @@ def _is_moe_experts(experts: Any) -> bool: def _shard_tensor_experts(experts: nn.Module, start: int, end: int) -> None: - experts.gate_up_proj = nn.Parameter(experts.gate_up_proj.data[start:end].clone()) - experts.down_proj = nn.Parameter(experts.down_proj.data[start:end].clone()) + experts.gate_up_proj = nn.Parameter( + experts.gate_up_proj.data[start:end].clone(), requires_grad=experts.gate_up_proj.requires_grad) + experts.down_proj = nn.Parameter( + experts.down_proj.data[start:end].clone(), requires_grad=experts.down_proj.requires_grad) if hasattr(experts, 'num_experts'): experts.num_experts = end - start @@ -441,9 +443,11 @@ def _shard_tensor_experts(experts: nn.Module, start: int, end: int) -> None: if not isinstance(target_param_wrapper, TargetParameterLoraWrapper): continue for tenant_name, tenant_tensor in target_param_wrapper.lora_A.items(): - target_param_wrapper.lora_A[tenant_name] = nn.Parameter(tenant_tensor.data[start:end].clone()) + target_param_wrapper.lora_A[tenant_name] = nn.Parameter( + tenant_tensor.data[start:end].clone(), requires_grad=tenant_tensor.requires_grad) for tenant_name, tenant_tensor in target_param_wrapper.lora_B.items(): - target_param_wrapper.lora_B[tenant_name] = nn.Parameter(tenant_tensor.data[start:end].clone()) + target_param_wrapper.lora_B[tenant_name] = nn.Parameter( + tenant_tensor.data[start:end].clone(), requires_grad=tenant_tensor.requires_grad) def _run_local_experts( diff --git a/src/twinkle/model/transformers/multi_lora_transformers.py b/src/twinkle/model/transformers/multi_lora_transformers.py index ea53930de..88b612223 100644 --- a/src/twinkle/model/transformers/multi_lora_transformers.py +++ b/src/twinkle/model/transformers/multi_lora_transformers.py @@ -81,10 +81,17 @@ def __init__( self.sp_strategy = None # Initialize expert parallel attributes (required by set_optimizer in TransformersModel) self.optimizer_group: Dict[str, OptimizerGroup] = {} - self.multi_adapter = MultiLora(max_loras=max_loras, max_r=max_r, max_length=max_length) + self.multi_adapter = MultiLora( + max_loras=max_loras, + max_r=max_r, + max_length=max_length, + defer_initial_weights=self._memory_efficient_init, + ) self.model.gradient_checkpointing_enable() self.model = self.multi_adapter.patch(self.model, target_modules=target_modules, lora_config=self.lora_config) - self.multi_adapter.save_initial_weights() + self._initial_lora_weights_saved = False + if not self._memory_efficient_init: + self._save_initial_lora_weights() # Active group for compatibility with single adapter self.active_group = None self.handler = self.register_global_mm_forward_hook() @@ -114,11 +121,18 @@ def unregister_mm_forward_hook(self, optimizer_group: OptimizerGroup): pass def _lazy_wrap_model(self): - return super()._lazy_wrap_model() + super()._lazy_wrap_model() + # Non-source ranks keep both PEFT and target-parameter LoRA slots on + # meta until NativeFSDPStrategy broadcasts the rank-local shards. + self._save_initial_lora_weights() + + def _save_initial_lora_weights(self): + if self._initial_lora_weights_saved: + return + self.multi_adapter.save_initial_weights() + self._initial_lora_weights_saved = True def _maybe_apply_expert_parallel(self): - if self._memory_efficient_init: - raise NotImplementedError('Expert parallel is not supported with memory_efficient_init') return super()._maybe_apply_expert_parallel() def _ensure_target_parameter_lora_installed(self, config: LoraConfig) -> None: @@ -127,9 +141,6 @@ def _ensure_target_parameter_lora_installed(self, config: LoraConfig) -> None: return if self._model_wrapped: raise RuntimeError('target_parameters LoRA must be installed before FSDP/DDP wrapping') - if getattr(self, '_enable_expert_parallel', False): - self.strategy.capture_pre_ep_state_if_needed(self.model, enable_ep=True) - # self._maybe_apply_expert_parallel() # 各rank广播之前不能对moe层进行分片, 没有实际权重时不能分片 self.multi_adapter.patch_target_parameters(self.model, target_parameters) @remote_function(dispatch='slice_dp', collect=collect_tensor_dict) @@ -280,6 +291,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) + self._lazy_wrap_model() with self.multi_adapter.save_context(kwargs.get('adapter_name')): load_optimizer = kwargs.get('load_optimizer', False) if output_dir is None: @@ -318,6 +330,7 @@ def calculate_metric(self, is_training, **kwargs): @remote_function() def remove_adapter(self, adapter_name: str): + self._lazy_wrap_model() if adapter_name in self.optimizer_group: self.optimizer_group.pop(adapter_name) self.multi_adapter.release_lora(adapter_name) diff --git a/src/twinkle/model/transformers/strategy/native_fsdp.py b/src/twinkle/model/transformers/strategy/native_fsdp.py index 925c84c82..97670494a 100644 --- a/src/twinkle/model/transformers/strategy/native_fsdp.py +++ b/src/twinkle/model/transformers/strategy/native_fsdp.py @@ -139,8 +139,14 @@ def wrap_model(self, model, optimizer=None): else: original_sd = model.state_dict() if is_source_rank else {} adapter_source_sd = _collect_adapter_source_state(model.state_dict()) - adapter_full_sd = ( - self._adapter_full_state_dict if is_source_rank and self._adapter_full_state_dict else {}) + if is_source_rank: + # Multi-LoRA target-parameter slots are installed before EP + # and then sharded with their experts. Preserve their full + # pre-EP tensors so every EP rank receives its own expert + # range instead of a copy of the source rank's local range. + adapter_full_sd = _collect_adapter_source_state(original_sd or {}, clone=False) + if self._adapter_full_state_dict: + adapter_full_sd.update(self._adapter_full_state_dict) saved_buffers = _get_non_persistent_buffers(model) if is_source_rank else {} if is_source_rank: model = model.to(torch.device('meta')) @@ -755,14 +761,15 @@ def _resolve_full_state_source_key(param_name: str, source_state: Mapping[str, A f'Tried source keys: {", ".join(candidates)}.') -def _collect_adapter_source_state(state_dict: Mapping[str, Any]) -> Dict[str, Any]: +def _collect_adapter_source_state(state_dict: Mapping[str, Any], *, clone: bool = True) -> Dict[str, Any]: adapter_state = {} for name, tensor in state_dict.items(): if not _is_lora_state_key(name) or not hasattr(tensor, 'detach'): continue if getattr(tensor, 'is_meta', False): continue - adapter_state[name] = tensor.detach().cpu().clone() + tensor = tensor.detach().cpu() + adapter_state[name] = tensor.clone() if clone else tensor return adapter_state diff --git a/tests/moe/test_ep_multi_lora_target_parameters.py b/tests/moe/test_ep_multi_lora_target_parameters.py index b92509d5f..90e721d40 100644 --- a/tests/moe/test_ep_multi_lora_target_parameters.py +++ b/tests/moe/test_ep_multi_lora_target_parameters.py @@ -2,6 +2,7 @@ import sys import torch import types +from torch import nn def _ensure_dummy_zmq(): @@ -24,7 +25,102 @@ def test_ep_target_parameter_lora_gather_dim_matches_peft_flattening(): assert _ep_expert_state_dict_gather_dim("model.layers.0.mlp.experts.base_layer.lora_A.weight") == 0 assert _ep_expert_state_dict_gather_dim("model.layers.0.mlp.experts.lora_B.weight") == 1 assert _ep_expert_state_dict_gather_dim("model.layers.0.mlp.experts.base_layer.lora_B.weight") == 1 - assert _ep_expert_state_dict_gather_dim("model.layers.0.mlp.experts._twinkle_lora_gate_up_proj.lora_B.lora_0.weight") == 0 + assert _ep_expert_state_dict_gather_dim( + "model.layers.0.mlp.experts._twinkle_lora_gate_up_proj.lora_B.lora_0.weight") == 0 + + +class _FakeTensorExperts(nn.Module): + + def __init__(self, *, device="cpu", dtype=torch.float32): + super().__init__() + self.gate_up_proj = nn.Parameter(torch.empty(4, 3, 8, device=device, dtype=dtype)) + self.down_proj = nn.Parameter(torch.empty(4, 4, 3, device=device, dtype=dtype)) + self.num_experts = 4 + + +def test_target_parameter_lora_slots_stay_meta_until_fsdp_materialization(): + from twinkle.model.multi_lora_target_parameters import TargetParameterLoraManager + + model = nn.Module() + model.experts = _FakeTensorExperts(device="meta") + manager = TargetParameterLoraManager(max_loras=2, max_r=4) + manager.patch(model, ["experts.gate_up_proj", "experts.down_proj"]) + + for wrapper in manager.wrappers: + assert all(param.is_meta for param in wrapper.lora_A.values()) + assert all(param.is_meta for param in wrapper.lora_B.values()) + + +def test_target_parameter_lora_defers_initial_snapshot_on_source_rank(): + from twinkle.model.multi_lora_target_parameters import TargetParameterLoraManager + + model = nn.Module() + model.experts = _FakeTensorExperts() + manager = TargetParameterLoraManager(max_loras=2, max_r=4, defer_initial_weights=True) + manager.patch(model, ["experts.gate_up_proj", "experts.down_proj"]) + + for wrapper in manager.wrappers: + assert all(not param.is_meta for param in wrapper.lora_A.values()) + assert all(torch.count_nonzero(param) == 0 for param in wrapper.lora_B.values()) + assert wrapper._initial_lora_A == {} + + manager.save_initial_weights() + + for wrapper in manager.wrappers: + assert set(wrapper._initial_lora_A) == {"lora_0", "lora_1"} + + +def test_ep_shards_target_parameter_lora_slots_on_meta(): + _ensure_dummy_zmq() + from twinkle.model.multi_lora_target_parameters import TargetParameterLoraManager + from twinkle.model.transformers.moe.expert_parallel import _shard_tensor_experts + + model = nn.Module() + model.experts = _FakeTensorExperts(device="meta") + manager = TargetParameterLoraManager(max_loras=2, max_r=4) + manager.patch(model, ["experts.gate_up_proj", "experts.down_proj"]) + + _shard_tensor_experts(model.experts, 2, 4) + + assert model.experts.gate_up_proj.shape[0] == 2 + assert model.experts.down_proj.shape[0] == 2 + for wrapper in manager.wrappers: + assert wrapper.num_experts == 2 + assert all(param.shape[0] == 2 and param.is_meta for param in wrapper.lora_A.values()) + assert all(param.shape[0] == 2 and param.is_meta for param in wrapper.lora_B.values()) + + +def test_target_parameter_slot_reset_uses_materialized_ep_local_snapshot(): + _ensure_dummy_zmq() + from peft import LoraConfig + from twinkle.model.multi_lora_target_parameters import TargetParameterLoraManager + from twinkle.model.transformers.moe.expert_parallel import _shard_tensor_experts + + torch.manual_seed(0) + model = nn.Module() + model.experts = _FakeTensorExperts() + manager = TargetParameterLoraManager(max_loras=1, max_r=4) + targets = ["experts.gate_up_proj", "experts.down_proj"] + manager.patch(model, targets) + manager.acquire( + "tenant_a", + "lora_0", + LoraConfig(r=2, lora_alpha=4, target_modules=[], target_parameters=targets), + ) + + _shard_tensor_experts(model.experts, 2, 4) + manager.save_initial_weights() + initial_a = [wrapper.lora_A["lora_0"].detach().clone() for wrapper in manager.wrappers] + + with torch.no_grad(): + for wrapper in manager.wrappers: + wrapper.lora_A["lora_0"].add_(1) + wrapper.lora_B["lora_0"].add_(1) + manager.release("tenant_a") + + for wrapper, expected_a in zip(manager.wrappers, initial_a): + assert torch.equal(wrapper.lora_A["lora_0"], expected_a) + assert torch.count_nonzero(wrapper.lora_B["lora_0"]) == 0 @pytest.mark.skipif(not torch.cuda.is_available() or torch.cuda.device_count() < 4, reason="Need 4 GPUs") From 72de402e53b2d19170790d91e0366b53ce94e1c8 Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Wed, 12 Aug 2026 17:41:55 +0800 Subject: [PATCH 02/28] fix(transformers): align multi-LoRA dtype before FSDP --- .../transformers/multi_lora_transformers.py | 7 +++ tests/model/test_multi_lora_dtype.py | 55 +++++++++++++++++++ 2 files changed, 62 insertions(+) create mode 100644 tests/model/test_multi_lora_dtype.py diff --git a/src/twinkle/model/transformers/multi_lora_transformers.py b/src/twinkle/model/transformers/multi_lora_transformers.py index 88b612223..a53d111eb 100644 --- a/src/twinkle/model/transformers/multi_lora_transformers.py +++ b/src/twinkle/model/transformers/multi_lora_transformers.py @@ -89,6 +89,13 @@ def __init__( ) self.model.gradient_checkpointing_enable() self.model = self.multi_adapter.patch(self.model, target_modules=target_modules, lora_config=self.lora_config) + # PEFT initializes LoRA parameters in FP32 even when the base model is + # BF16. Native FSDP2 records the pre-wrap parameter dtype as the + # DTensor grad_dtype, so materializing those slots from the BF16 rank-0 + # state later would make FP32 reduced gradients incompatible with the + # BF16 sharded parameters. Keep all preallocated slots aligned before + # EP state capture and FSDP wrapping on every rank, including meta ranks. + self._ensure_lora_dtype(self.model) self._initial_lora_weights_saved = False if not self._memory_efficient_init: self._save_initial_lora_weights() diff --git a/tests/model/test_multi_lora_dtype.py b/tests/model/test_multi_lora_dtype.py new file mode 100644 index 000000000..1a2a78102 --- /dev/null +++ b/tests/model/test_multi_lora_dtype.py @@ -0,0 +1,55 @@ +import sys +import torch +import types +from peft import LoraConfig +from torch import nn + + +def _ensure_dummy_zmq(): + if 'zmq' in sys.modules: + return + sys.modules['zmq'] = types.SimpleNamespace( + Context=object, + Socket=object, + RCVTIMEO=1, + SNDTIMEO=2, + LINGER=3, + ) + + +def _make_multi_lora_model(device: str): + _ensure_dummy_zmq() + from twinkle.model.multi_lora import MultiLora + + model = nn.Sequential(nn.Linear(4, 4, device=device, dtype=torch.bfloat16)) + multi_lora = MultiLora(max_loras=2, max_r=4, defer_initial_weights=True) + model = multi_lora.patch( + model, + target_modules=['0'], + lora_config=LoraConfig(r=4, lora_alpha=8, target_modules=['0']), + ) + return model + + +def _align_lora_dtype(model): + _ensure_dummy_zmq() + from twinkle.model.transformers.transformers import TransformersModel + + TransformersModel._ensure_lora_dtype(None, model) + + +def test_multi_lora_dtype_matches_bf16_base_before_fsdp_wrap(): + model = _make_multi_lora_model('cpu') + + assert {param.dtype for name, param in model.named_parameters() if 'lora_' in name} == {torch.float32} + _align_lora_dtype(model) + assert {param.dtype for name, param in model.named_parameters() if 'lora_' in name} == {torch.bfloat16} + + +def test_meta_multi_lora_dtype_matches_bf16_base_before_fsdp_wrap(): + model = _make_multi_lora_model('meta') + + assert all(param.is_meta for param in model.parameters()) + assert {param.dtype for name, param in model.named_parameters() if 'lora_' in name} == {torch.float32} + _align_lora_dtype(model) + assert {param.dtype for name, param in model.named_parameters() if 'lora_' in name} == {torch.bfloat16} From 6f339e05b6d8602c76edd3a6103701e64a036e3f Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Wed, 12 Aug 2026 20:23:49 +0800 Subject: [PATCH 03/28] wip --- .../transformer/server_config_dsv4_0731.yaml | 120 ++++++++++++++++ .../server/transformer/server_dsv4_0731.py | 12 ++ .../client/twinkle/dsv4_multi_lora_sft.py | 132 ++++++++++++++++++ read_me.md | 56 ++++++++ .../model/multi_lora_target_parameters.py | 4 + .../transformers/multi_lora_transformers.py | 11 ++ src/twinkle/server/config/application_spec.py | 8 ++ .../test_ep_multi_lora_target_parameters.py | 19 +++ tests/server/config/test_server_config.py | 1 + 9 files changed, 363 insertions(+) create mode 100644 cookbook/client/server/transformer/server_config_dsv4_0731.yaml create mode 100644 cookbook/client/server/transformer/server_dsv4_0731.py create mode 100644 cookbook/client/twinkle/dsv4_multi_lora_sft.py create mode 100644 read_me.md diff --git a/cookbook/client/server/transformer/server_config_dsv4_0731.yaml b/cookbook/client/server/transformer/server_config_dsv4_0731.yaml new file mode 100644 index 000000000..2b93293ed --- /dev/null +++ b/cookbook/client/server/transformer/server_config_dsv4_0731.yaml @@ -0,0 +1,120 @@ +# DeepSeek-V4-Flash-0731 client/server Multi-LoRA SFT on one 2-GPU node. +# Both GPUs are assigned to the Transformers training model. A vLLM sampler +# is intentionally omitted because it would compete with training for GPUs. + +proxy_location: EveryNode + +http_options: + host: 0.0.0.0 + port: 8000 + +telemetry: + enabled: false + otlp_endpoint: http://localhost:4317 + +persistence: + mode: file + file_path: /tmp/twinkle_dsv4_0731_state.json + +applications: + - name: server + route_prefix: /api/v1 + import_path: server + args: + server_config: + per_token_model_limit: 2 + # This public name is used by the client and HTTP routes. The launcher + # maps it to the local model_id configured on the model deployment. + supported_models: + - deepseek-v4-0731-local + deployments: + - name: TinkerCompatServer + max_ongoing_requests: 50 + autoscaling_config: + min_replicas: 1 + max_replicas: 1 + target_ongoing_requests: 32 + ray_actor_options: + num_cpus: 0.1 + runtime_env: + env_vars: + TWINKLE_FAIL_FAST: "1" + + - name: models-deepseek-v4-0731-local + route_prefix: /api/v1/model/deepseek-v4-0731-local + import_path: model + args: + backend: transformers + model_id: /nas/disk1/random-deepseek-v4-4b + nproc_per_node: 2 + max_loras: 2 + max_r: 8 + max_length: 2048 + mixed_precision: bf16 + strategy: native_fsdp + memory_efficient_init: true + target_modules: all-linear + fsdp_config: + reshard_after_forward: true + expert_parallel: + enabled: true + ep_size: 2 + router_dtype: fp32 + keep_router_logits: false + device_group: + name: model + ranks: 2 + device_type: cuda + device_mesh: + device_type: cuda + fsdp_size: 2 + dp_size: 1 + ep_size: 2 + queue_config: + rps_limit: 20 + tps_limit: 100000 + max_input_tokens: 4096 + queue_timeout: 3600 + execution_timeout: 3600 + adapter_config: + adapter_timeout: 3600 + adapter_max_lifetime: 86400 + deployments: + - name: ModelManagement + autoscaling_config: + min_replicas: 1 + max_replicas: 1 + target_ongoing_requests: 8 + ray_actor_options: + num_cpus: 0.1 + runtime_env: + env_vars: + TWINKLE_TRUST_REMOTE_CODE: "1" + TWINKLE_FAIL_FAST: "1" + + # Dataset/tokenization stays on CPU and reads the local JSONL path from the + # machine running the server. + - name: processor + route_prefix: /api/v1/processor + import_path: processor + args: + ncpu_proc_per_node: 2 + device_group: + name: processor + ranks: 2 + device_type: CPU + device_mesh: + device_type: CPU + dp_size: 2 + deployments: + - name: ProcessorManagement + autoscaling_config: + min_replicas: 1 + max_replicas: 1 + target_ongoing_requests: 32 + ray_actor_options: + num_cpus: 0.1 + runtime_env: + env_vars: + TWINKLE_TRUST_REMOTE_CODE: "1" + TWINKLE_FAIL_FAST: "1" diff --git a/cookbook/client/server/transformer/server_dsv4_0731.py b/cookbook/client/server/transformer/server_dsv4_0731.py new file mode 100644 index 000000000..3e8de2797 --- /dev/null +++ b/cookbook/client/server/transformer/server_dsv4_0731.py @@ -0,0 +1,12 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Launch the two-GPU DeepSeek-V4-0731 Multi-LoRA server.""" +import os + +os.environ.setdefault('TWINKLE_TRUST_REMOTE_CODE', '1') + +from twinkle.server import launch_server # noqa: E402 + +file_dir = os.path.abspath(os.path.dirname(__file__)) +config_path = os.path.join(file_dir, 'server_config_dsv4_0731.yaml') + +launch_server(config_path=config_path) diff --git a/cookbook/client/twinkle/dsv4_multi_lora_sft.py b/cookbook/client/twinkle/dsv4_multi_lora_sft.py new file mode 100644 index 000000000..41f58dfcd --- /dev/null +++ b/cookbook/client/twinkle/dsv4_multi_lora_sft.py @@ -0,0 +1,132 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Client-side Multi-LoRA SFT for a local DeepSeek-V4-Flash-0731 server. + +Start ``server_config_dsv4_0731.yaml`` first. This client creates two LoRA +tenants on the same two-GPU EP/FSDP model and alternates their SFT micro-steps. +The base model and dataset are both local paths on the server machine. +""" +import os + +from peft import LoraConfig + +from twinkle import get_logger, init_twinkle_client +from twinkle.dataset import DatasetMeta +from twinkle_client.dataloader import DataLoader +from twinkle_client.dataset import Dataset +from twinkle_client.model import MultiLoraTransformersModel + +logger = get_logger() + +SERVER_URL = os.environ.get('TWINKLE_SERVER_URL', 'http://localhost:8000') +SERVER_TOKEN = os.environ.get('TWINKLE_SERVER_TOKEN', 'EMPTY_TOKEN') +SERVED_MODEL_NAME = os.environ.get('TWINKLE_MODEL_ID', 'deepseek-v4-0731-local') +MODEL_PATH = os.environ.get('DSV4_MODEL_ID', '/nas/disk1/random-deepseek-v4-4b') +DATASET_PATH = os.environ.get('DATASET_ID', '/model/ljl/dataset/self-cognition.jsonl') +OUTPUT_DIR = os.environ.get('OUTPUT_DIR', '/tmp/twinkle_dsv4_0731_multi_lora') + +ADAPTER_NAMES = tuple( + name.strip() for name in os.environ.get('ADAPTER_NAMES', 'tenant_a,tenant_b').split(',') if name.strip()) +BATCH_SIZE = int(os.environ.get('BATCH_SIZE', '2')) +GRAD_ACCUM_STEPS = int(os.environ.get('GRAD_ACCUM_STEPS', '4')) +MAX_STEPS = int(os.environ.get('MAX_STEPS', '10')) +MAX_LENGTH = int(os.environ.get('MAX_LENGTH', '2048')) +LR = float(os.environ.get('LR', '1e-4')) +LORA_R = int(os.environ.get('LORA_R', '8')) +LORA_ALPHA = int(os.environ.get('LORA_ALPHA', '32')) + + +def _build_dataset() -> Dataset: + dataset = Dataset(dataset_meta=DatasetMeta(DATASET_PATH)) + dataset.set_template('DeepseekV4Template', model_id=MODEL_PATH, max_length=MAX_LENGTH) + dataset.map( + 'SelfCognitionProcessor', + init_args={ + 'model_name': 'twinkle模型', + 'model_author': 'ModelScope社区', + }, + ) + dataset.encode() + return dataset + + +def _build_lora_config() -> LoraConfig: + return LoraConfig( + r=LORA_R, + lora_alpha=LORA_ALPHA, + lora_dropout=0.0, + target_modules='all-linear', + exclude_modules=['o_a_proj'], + target_parameters=[ + 'mlp.experts.gate_up_proj', + 'mlp.experts.down_proj', + ], + ) + + +def train() -> None: + client = init_twinkle_client(base_url=SERVER_URL, api_key=SERVER_TOKEN) + supported_models = [item.model_name for item in client.get_server_capabilities().supported_models] + if SERVED_MODEL_NAME not in supported_models: + raise RuntimeError(f'{SERVED_MODEL_NAME!r} is not served; available models: {supported_models}') + if len(ADAPTER_NAMES) != 2: + raise ValueError('This two-slot server example requires exactly two ADAPTER_NAMES.') + if BATCH_SIZE < 2: + raise ValueError('BATCH_SIZE must be at least 2 because the model uses two FSDP data ranks.') + + dataset = _build_dataset() + dataloaders = {name: DataLoader(dataset=dataset, batch_size=BATCH_SIZE) for name in ADAPTER_NAMES} + models = {name: MultiLoraTransformersModel(model_id=SERVED_MODEL_NAME) for name in ADAPTER_NAMES} + lora_config = _build_lora_config() + + # Register every tenant before set_optimizer(). The first optimizer call + # materializes and EP/FSDP-shards all preallocated Multi-LoRA slots. + for name, model in models.items(): + save_dir = os.path.join(OUTPUT_DIR, name) + os.makedirs(save_dir, exist_ok=True) + model.add_adapter_to_model( + name, + lora_config, + gradient_accumulation_steps=GRAD_ACCUM_STEPS, + save_dir=save_dir, + ) + model.set_template('DeepseekV4Template', max_length=MAX_LENGTH) + model.set_processor('InputProcessor', padding_side='right') + model.set_loss('CrossEntropyLoss') + + for model in models.values(): + model.set_optimizer('AdamW', lr=LR, foreach=False) + + iterators = {name: iter(loader) for name, loader in dataloaders.items()} + completed_steps = {name: 0 for name in ADAPTER_NAMES} + active_names = set(ADAPTER_NAMES) + + while active_names and any(step < MAX_STEPS for step in completed_steps.values()): + for name in ADAPTER_NAMES: + if name not in active_names or completed_steps[name] >= MAX_STEPS: + continue + try: + batch = next(iterators[name]) + except StopIteration: + active_names.remove(name) + continue + + model = models[name] + model.forward_backward(inputs=batch, gradient_accumulation_steps=GRAD_ACCUM_STEPS) + model.clip_grad_and_step(max_grad_norm=1.0, gradient_accumulation_steps=GRAD_ACCUM_STEPS) + completed_steps[name] += 1 + + if completed_steps[name] % GRAD_ACCUM_STEPS == 0: + metric = model.calculate_metric(is_training=True) + logger.info('adapter=%s micro_step=%s metric=%s', name, completed_steps[name], metric.result) + + for name, model in models.items(): + checkpoint = model.save( + name=f'dsv4-0731-{name}-final', + save_optimizer=True, + consumed_train_samples=dataloaders[name].get_state()['consumed_train_samples'], + ) + logger.info('Saved adapter %s: %s', name, checkpoint.twinkle_path) + + +if __name__ == '__main__': + train() diff --git a/read_me.md b/read_me.md new file mode 100644 index 000000000..295b4c3e4 --- /dev/null +++ b/read_me.md @@ -0,0 +1,56 @@ +已写好。既然模型目录存在 `tokenizer.json`，会直接从 `/nas/disk1/random-deepseek-v4-4b` 加载，不需要单独配置 `tokenizer_id`。 + +生成的文件： + +- [server_config_dsv4_0731.yaml](/Users/linjiajia/project/twinkle/cookbook/client/server/transformer/server_config_dsv4_0731.yaml) +- [server_dsv4_0731.py](/Users/linjiajia/project/twinkle/cookbook/client/server/transformer/server_dsv4_0731.py) +- [dsv4_multi_lora_sft.py](/Users/linjiajia/project/twinkle/cookbook/client/twinkle/dsv4_multi_lora_sft.py) + +配置为： + +- 本地模型：`/nas/disk1/random-deepseek-v4-4b` +- 本地数据集：`/model/ljl/dataset/self-cognition.jsonl` +- 2 张 GPU +- Native FSDP2 + EP=2 +- `memory_efficient_init: true` +- 两个 LoRA：`tenant_a`、`tenant_b` +- `DeepseekV4Template` +- 两张 GPU 全部用于训练，因此没有启动 vLLM sampler + +服务端启动： + +```bash +cd /model/ljl/project/remote-git/dsv4_0731/twinkle +export PYTHONPATH="$PWD/src:$PYTHONPATH" + +CUDA_VISIBLE_DEVICES=0,1 ray start \ + --head \ + --num-gpus=2 \ + --disable-usage-stats \ + --include-dashboard=false + +CUDA_VISIBLE_DEVICES=0,1 python3 \ + cookbook/client/server/transformer/server_dsv4_0731.py +``` + +如果 Ray 集群已经启动，跳过 `ray start`。 + +另一个终端启动客户端： + +```bash +cd /model/ljl/project/remote-git/dsv4_0731/twinkle +export PYTHONPATH="$PWD/src:$PYTHONPATH" + +TWINKLE_SERVER_URL=http://127.0.0.1:8000 \ +TWINKLE_SERVER_TOKEN=EMPTY_TOKEN \ +python3 cookbook/client/twinkle/dsv4_multi_lora_sft.py +``` + +`DeepseekV4Template` 会优先加载模型目录里的： + +```text +tokenizer.json +encoding/encoding_dsv4.py +``` + +若不存在 `encoding/encoding_dsv4.py`，才回退到 Twinkle 内置 encoding。相关检查共 `31 passed, 1 skipped`；当前改动尚未 commit/push。 \ No newline at end of file diff --git a/src/twinkle/model/multi_lora_target_parameters.py b/src/twinkle/model/multi_lora_target_parameters.py index d6936776c..b8af6c4d2 100644 --- a/src/twinkle/model/multi_lora_target_parameters.py +++ b/src/twinkle/model/multi_lora_target_parameters.py @@ -339,6 +339,10 @@ def __init__(self, max_loras: int, max_r: int, defer_initial_weights: bool = Fal self.tenant_configs: dict[str, LoraConfig] = {} self._target_parameters: tuple[str, ...] | None = None + @property + def patched_target_parameters(self) -> tuple[str, ...] | None: + return self._target_parameters + def patch(self, model: nn.Module, target_parameters: Iterable[str]) -> None: target_parameters = tuple(target_parameters) if not target_parameters: diff --git a/src/twinkle/model/transformers/multi_lora_transformers.py b/src/twinkle/model/transformers/multi_lora_transformers.py index a53d111eb..0f95d937d 100644 --- a/src/twinkle/model/transformers/multi_lora_transformers.py +++ b/src/twinkle/model/transformers/multi_lora_transformers.py @@ -146,6 +146,13 @@ def _ensure_target_parameter_lora_installed(self, config: LoraConfig) -> None: target_parameters = getattr(config, 'target_parameters', None) if not target_parameters: return + target_parameter_manager = self.multi_adapter.target_parameter_manager + if target_parameter_manager.patched_target_parameters is not None: + # The first tenant preallocates target-parameter slots for every + # LoRA. Later tenants can reuse the same slots after EP/FSDP wrap; + # patch() still rejects a different target set. + self.multi_adapter.patch_target_parameters(self.model, target_parameters) + return if self._model_wrapped: raise RuntimeError('target_parameters LoRA must be installed before FSDP/DDP wrapping') self.multi_adapter.patch_target_parameters(self.model, target_parameters) @@ -231,6 +238,10 @@ def set_loss(self, loss_cls: Union[Type[Loss], str], **kwargs): @remote_function() def set_optimizer(self, optimizer_cls: Union[Type[Optimizer], str], **kwargs): self._check_adapter_valid(kwargs.get('adapter_name')) + # Materialize/shard the preallocated LoRA slots before an optimizer + # captures parameter references. Otherwise the first optimizer can + # retain pre-FSDP/meta parameters while later optimizers see DTensors. + self._lazy_wrap_model() with self.multi_adapter.adapter(kwargs.get('adapter_name')): super().set_optimizer(optimizer_cls, **kwargs) diff --git a/src/twinkle/server/config/application_spec.py b/src/twinkle/server/config/application_spec.py index 3f8d60cbc..4be21d746 100644 --- a/src/twinkle/server/config/application_spec.py +++ b/src/twinkle/server/config/application_spec.py @@ -58,7 +58,15 @@ class ModelArgs(_ArgsBase): adapter_config: dict[str, Any] | None = None queue_config: TaskQueueConfig = Field(default_factory=TaskQueueConfig) max_loras: int = 5 + max_r: int | None = None max_length: int | None = None + mixed_precision: Literal['no', 'fp8', 'fp16', 'bf16'] | None = None + strategy: Literal['accelerate', 'native_fsdp'] | None = None + ddp_config: dict[str, Any] | None = None + fsdp_config: dict[str, Any] | None = None + grad_scaler_config: dict[str, Any] | None = None + memory_efficient_init: bool | None = None + target_modules: str | list[str] | None = None class SamplerArgs(_ArgsBase): diff --git a/tests/moe/test_ep_multi_lora_target_parameters.py b/tests/moe/test_ep_multi_lora_target_parameters.py index 90e721d40..3266e2b2d 100644 --- a/tests/moe/test_ep_multi_lora_target_parameters.py +++ b/tests/moe/test_ep_multi_lora_target_parameters.py @@ -70,6 +70,25 @@ def test_target_parameter_lora_defers_initial_snapshot_on_source_rank(): assert set(wrapper._initial_lora_A) == {"lora_0", "lora_1"} +def test_target_parameter_lora_reuses_matching_preallocated_slots(): + from twinkle.model.multi_lora_target_parameters import TargetParameterLoraManager + + model = nn.Module() + model.experts = _FakeTensorExperts() + manager = TargetParameterLoraManager(max_loras=2, max_r=4) + targets = ["experts.gate_up_proj", "experts.down_proj"] + + manager.patch(model, targets) + wrappers = list(manager.wrappers) + manager.patch(model, targets) + + assert manager.patched_target_parameters == tuple(targets) + assert manager.wrappers == wrappers + + with pytest.raises(ValueError, match="target_parameters already patched"): + manager.patch(model, ["experts.gate_up_proj"]) + + def test_ep_shards_target_parameter_lora_slots_on_meta(): _ensure_dummy_zmq() from twinkle.model.multi_lora_target_parameters import TargetParameterLoraManager diff --git a/tests/server/config/test_server_config.py b/tests/server/config/test_server_config.py index 93644cbef..ca62716d0 100644 --- a/tests/server/config/test_server_config.py +++ b/tests/server/config/test_server_config.py @@ -240,6 +240,7 @@ def test_cookbook_examples_load() -> None: here = Path(__file__).resolve().parents[3] examples = [ here / 'cookbook' / 'client' / 'server' / 'transformer' / 'server_config.yaml', + here / 'cookbook' / 'client' / 'server' / 'transformer' / 'server_config_dsv4_0731.yaml', here / 'cookbook' / 'client' / 'server' / 'megatron' / 'server_config.yaml', here / 'cookbook' / 'client' / 'server' / 'megatron' / 'server_config_4b.yaml', ] From 3893b1cbf2c3016535ecec806a148c64c5660770 Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Wed, 12 Aug 2026 20:51:21 +0800 Subject: [PATCH 04/28] wip --- src/twinkle/infra/_ray/ray_helper.py | 31 ++++++++++ src/twinkle/infra/_ray/resource_manager.py | 11 +++- .../transformers/strategy/native_fsdp.py | 53 +++++++++++++---- .../model/transformers/transformers.py | 3 + tests/infra/test_ray_worker_topology.py | 27 +++++++++ .../test_native_fsdp_node_topology.py | 58 +++++++++++++++++++ 6 files changed, 169 insertions(+), 14 deletions(-) create mode 100644 tests/infra/test_ray_worker_topology.py create mode 100644 tests/transformers/test_native_fsdp_node_topology.py diff --git a/src/twinkle/infra/_ray/ray_helper.py b/src/twinkle/infra/_ray/ray_helper.py index 6eb991383..c3792812b 100644 --- a/src/twinkle/infra/_ray/ray_helper.py +++ b/src/twinkle/infra/_ray/ray_helper.py @@ -8,6 +8,21 @@ T = TypeVar('T') +def _get_node_local_topology(placements: List[Dict[str, Any]]) -> List[Tuple[int, List[int]]]: + """Return each worker's node-local index and distributed ranks on that node.""" + node_to_ranks: Dict[int, List[int]] = {} + for rank, placement in enumerate(placements): + node_rank = int(placement.get('node_rank', 0)) + node_to_ranks.setdefault(node_rank, []).append(rank) + + topology = [] + for rank, placement in enumerate(placements): + node_rank = int(placement.get('node_rank', 0)) + node_ranks = node_to_ranks[node_rank] + topology.append((node_ranks.index(rank), node_ranks)) + return topology + + class RayHelper: resource_manager: Optional[ResourceManager] = None @@ -304,6 +319,7 @@ def create_workers(worker_cls: Type[T], ip, port = RayHelper.get_master_id_port(placement_groups[0]['placement_group']) device_type_upper = (device_config.device_type or '').upper() + node_local_topology = _get_node_local_topology(placement_groups) if device_type_upper != 'CPU': world_size = len(ranks) device_type = Platform.get_platform(device_type_upper).__name__ @@ -319,6 +335,17 @@ def create_workers(worker_cls: Type[T], str(pg_idx), 'LOCAL_RANK': str(0), + # Each Ray actor sees only its own accelerator, so LOCAL_RANK + # must remain 0 as the device index. Keep the node-local + # distributed topology in separate Twinkle variables. + 'TWINKLE_NODE_LOCAL_RANK': + str(node_local_topology[pg_idx][0]), + 'TWINKLE_NODE_LOCAL_WORLD_SIZE': + str(len(node_local_topology[pg_idx][1])), + 'TWINKLE_NODE_RANKS': + ','.join(str(rank) for rank in node_local_topology[pg_idx][1]), + 'NODE_RANK': + str(deploy_pg.get('node_rank', 0)), 'CLUSTER_NAME': cluster_name, 'WORKER_NAME': @@ -373,6 +400,10 @@ def create_workers(worker_cls: Type[T], 'WORLD_SIZE': str(world_size), 'RANK': str(rank), 'LOCAL_RANK': str(0), + 'TWINKLE_NODE_LOCAL_RANK': str(node_local_topology[rank][0]), + 'TWINKLE_NODE_LOCAL_WORLD_SIZE': str(len(node_local_topology[rank][1])), + 'TWINKLE_NODE_RANKS': ','.join(str(item) for item in node_local_topology[rank][1]), + 'NODE_RANK': str(deploy_pg.get('node_rank', 0)), 'CLUSTER_NAME': cluster_name, 'WORKER_NAME': worker_name, 'TWINKLE_MODE': 'ray', diff --git a/src/twinkle/infra/_ray/resource_manager.py b/src/twinkle/infra/_ray/resource_manager.py index e9cba0f4d..c51e0a512 100644 --- a/src/twinkle/infra/_ray/resource_manager.py +++ b/src/twinkle/infra/_ray/resource_manager.py @@ -217,6 +217,7 @@ def get_visible_devices(): local_device_groups.append( dict( gpu_rank=gpu_ranks_local, + node_rank=node_rank, placement_group=self.node2pg[node_rank], ray_address=ray_address)) else: @@ -224,7 +225,11 @@ def get_visible_devices(): node_rank = alloc_rank // nproc_per_node gpu_rank = self.visible_devices[node_rank - self.min_node_idx][alloc_rank % nproc_per_node] local_device_groups.append( - dict(gpu_rank=[gpu_rank], placement_group=self.node2pg[node_rank], ray_address=ray_address)) + dict( + gpu_rank=[gpu_rank], + node_rank=node_rank, + placement_group=self.node2pg[node_rank], + ray_address=ray_address)) self.device_groups[group.name] = local_device_groups @@ -237,9 +242,11 @@ def get_visible_devices(): ranks = group.ranks local_device_groups = [] for _ in range(ranks): + node_rank = self.cpu_node_map[global_cpu_proc_idx][0] local_device_groups.append( dict( - placement_group=self.cpu_placement_groups[self.cpu_node_map[global_cpu_proc_idx][0]], + node_rank=node_rank, + placement_group=self.cpu_placement_groups[node_rank], ray_address=ray_address)) global_cpu_proc_idx += 1 self.device_groups[group.name] = local_device_groups diff --git a/src/twinkle/model/transformers/strategy/native_fsdp.py b/src/twinkle/model/transformers/strategy/native_fsdp.py index 97670494a..70b9fd5b3 100644 --- a/src/twinkle/model/transformers/strategy/native_fsdp.py +++ b/src/twinkle/model/transformers/strategy/native_fsdp.py @@ -19,6 +19,18 @@ LORA_STATE_KEY_MARKERS = ('lora_A', 'lora_B', 'lora_embedding') PEFT_BASE_PREFIX = 'base_model.model.' PEFT_BASE_LAYER_SEGMENT = 'base_layer' +TWINKLE_NODE_LOCAL_RANK = 'TWINKLE_NODE_LOCAL_RANK' +TWINKLE_NODE_LOCAL_WORLD_SIZE = 'TWINKLE_NODE_LOCAL_WORLD_SIZE' +TWINKLE_NODE_RANKS = 'TWINKLE_NODE_RANKS' + + +def _get_node_local_rank() -> int: + """Return process topology rank without changing the actor-local device index.""" + return int(os.environ.get(TWINKLE_NODE_LOCAL_RANK, Platform.get_local_rank())) + + +def _get_node_local_world_size() -> int: + return int(os.environ.get(TWINKLE_NODE_LOCAL_WORLD_SIZE, Platform.get_local_world_size())) class NativeFSDPStrategy: @@ -49,14 +61,20 @@ def pretrained_load_context(self): def use_rank0_pretrained_broadcast(self) -> bool: return self._memory_efficient_init and self.device_mesh is not None + def is_node_local_source_rank(self) -> bool: + local_rank = _get_node_local_rank() + if local_rank < 0: + raise RuntimeError('Native FSDP memory_efficient_init requires node-local rank topology.') + return local_rank == 0 + def capture_pre_ep_state_if_needed(self, model, *, enable_ep: bool) -> None: if self._pre_ep_state_captured: return if not (enable_ep and self.use_rank0_pretrained_broadcast()): return - local_rank = Platform.get_local_rank() + local_rank = _get_node_local_rank() if local_rank < 0: - raise RuntimeError('Native FSDP node-local pre-EP state capture requires LOCAL_RANK.') + raise RuntimeError('Native FSDP node-local pre-EP state capture requires node-local rank topology.') is_source_rank = dist.is_available() and dist.is_initialized() and local_rank == 0 self.set_rank0_pre_ep_full_state_dict(clone_state_dict_to_cpu(model.state_dict()) if is_source_rank else {}) self._pre_ep_state_captured = True @@ -130,9 +148,9 @@ def wrap_model(self, model, optimizer=None): adapter_source_sd = {} adapter_full_sd = {} if use_meta: - local_rank = Platform.get_local_rank() + local_rank = _get_node_local_rank() if local_rank < 0: - raise RuntimeError('Native FSDP node-local state loading requires LOCAL_RANK.') + raise RuntimeError('Native FSDP node-local state loading requires node-local rank topology.') is_source_rank = local_rank == 0 if ep_enabled and self._rank0_pre_ep_full_state_dict is not None: original_sd = self._rank0_pre_ep_full_state_dict if is_source_rank else {} @@ -600,19 +618,30 @@ def _get_local_rank_info() -> tuple[int, int, int, List[int]]: """Return local-rank topology for node-local state-dict fanout.""" rank = dist.get_rank() world_size = dist.get_world_size() - local_rank = Platform.get_local_rank() - if 'LOCAL_WORLD_SIZE' not in os.environ and 'LOCAL_SIZE' not in os.environ: - raise RuntimeError('Native FSDP node-local state loading requires LOCAL_WORLD_SIZE or LOCAL_SIZE.') - local_world_size = Platform.get_local_world_size() - if local_rank < 0 or local_world_size <= 0 or world_size % local_world_size != 0: + local_rank = _get_node_local_rank() + has_twinkle_topology = TWINKLE_NODE_LOCAL_WORLD_SIZE in os.environ + if not has_twinkle_topology and 'LOCAL_WORLD_SIZE' not in os.environ and 'LOCAL_SIZE' not in os.environ: + raise RuntimeError( + 'Native FSDP node-local state loading requires Twinkle Ray worker topology, ' + 'LOCAL_WORLD_SIZE, or LOCAL_SIZE.') + local_world_size = _get_node_local_world_size() + raw_node_ranks = os.environ.get(TWINKLE_NODE_RANKS) + if (local_rank < 0 or local_world_size <= 0 + or (not raw_node_ranks and world_size % local_world_size != 0)): raise RuntimeError(f'Invalid local rank topology: rank={rank}, world_size={world_size}, ' f'local_rank={local_rank}, local_world_size={local_world_size}.') - node_start = rank - local_rank - node_ranks = list(range(node_start, min(node_start + local_world_size, world_size))) + if raw_node_ranks: + node_ranks = [int(item) for item in raw_node_ranks.split(',') if item] + else: + node_start = rank - local_rank + node_ranks = list(range(node_start, min(node_start + local_world_size, world_size))) if rank not in node_ranks or len(node_ranks) != local_world_size: raise RuntimeError(f'Invalid local rank group: rank={rank}, local_rank={local_rank}, ' f'local_world_size={local_world_size}, node_ranks={node_ranks}.') - return rank, world_size, node_start, node_ranks + if node_ranks[local_rank] != rank: + raise RuntimeError(f'Invalid node-local rank ordering: rank={rank}, local_rank={local_rank}, ' + f'node_ranks={node_ranks}.') + return rank, world_size, node_ranks[0], node_ranks def _find_experts_in_layer(layer_mod: nn.Module, experts_map: Dict[str, nn.Module]) -> Optional[nn.Module]: diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index b561c1664..5f7eb251f 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -314,6 +314,9 @@ def _should_init_empty_pretrained_model_on_this_rank(self) -> bool: use_rank0_broadcast = getattr(self.strategy, 'use_rank0_pretrained_broadcast', lambda: False) if not (use_rank0_broadcast() and dist.is_available() and dist.is_initialized()): return False + is_node_local_source_rank = getattr(self.strategy, 'is_node_local_source_rank', None) + if is_node_local_source_rank is not None: + return not is_node_local_source_rank() local_rank = Platform.get_local_rank() if local_rank < 0: raise RuntimeError('Native FSDP memory_efficient_init requires LOCAL_RANK.') diff --git a/tests/infra/test_ray_worker_topology.py b/tests/infra/test_ray_worker_topology.py new file mode 100644 index 000000000..a4098bd4e --- /dev/null +++ b/tests/infra/test_ray_worker_topology.py @@ -0,0 +1,27 @@ +from twinkle.infra._ray.ray_helper import _get_node_local_topology + + +def test_get_node_local_topology_for_single_node_workers(): + placements = [ + {'node_rank': 0}, + {'node_rank': 0}, + ] + + assert _get_node_local_topology(placements) == [ + (0, [0, 1]), + (1, [0, 1]), + ] + + +def test_get_node_local_topology_does_not_assume_contiguous_global_ranks(): + placements = [ + {'node_rank': 0}, + {'node_rank': 1}, + {'node_rank': 0}, + ] + + assert _get_node_local_topology(placements) == [ + (0, [0, 2]), + (0, [1]), + (1, [0, 2]), + ] diff --git a/tests/transformers/test_native_fsdp_node_topology.py b/tests/transformers/test_native_fsdp_node_topology.py new file mode 100644 index 000000000..b22f9654b --- /dev/null +++ b/tests/transformers/test_native_fsdp_node_topology.py @@ -0,0 +1,58 @@ +import torch.distributed as dist + +from twinkle.model.transformers.strategy.native_fsdp import ( + NativeFSDPStrategy, + _get_local_rank_info, +) + + +_TOPOLOGY_ENV_NAMES = ( + 'TWINKLE_NODE_LOCAL_RANK', + 'TWINKLE_NODE_LOCAL_WORLD_SIZE', + 'TWINKLE_NODE_RANKS', + 'LOCAL_RANK', + 'LOCAL_WORLD_SIZE', + 'LOCAL_SIZE', +) + + +def _clear_topology_env(monkeypatch): + for name in _TOPOLOGY_ENV_NAMES: + monkeypatch.delenv(name, raising=False) + + +def test_native_fsdp_uses_ray_node_rank_instead_of_actor_device_index(monkeypatch): + _clear_topology_env(monkeypatch) + monkeypatch.setenv('LOCAL_RANK', '0') + monkeypatch.setenv('TWINKLE_NODE_LOCAL_RANK', '1') + monkeypatch.setenv('TWINKLE_NODE_LOCAL_WORLD_SIZE', '2') + monkeypatch.setenv('TWINKLE_NODE_RANKS', '0,1') + monkeypatch.setattr(dist, 'get_rank', lambda: 1) + monkeypatch.setattr(dist, 'get_world_size', lambda: 2) + + strategy = NativeFSDPStrategy(device_mesh=None, memory_efficient_init=True) + + assert strategy.is_node_local_source_rank() is False + assert _get_local_rank_info() == (1, 2, 0, [0, 1]) + + +def test_native_fsdp_accepts_explicit_noncontiguous_node_ranks(monkeypatch): + _clear_topology_env(monkeypatch) + monkeypatch.setenv('LOCAL_RANK', '0') + monkeypatch.setenv('TWINKLE_NODE_LOCAL_RANK', '1') + monkeypatch.setenv('TWINKLE_NODE_LOCAL_WORLD_SIZE', '2') + monkeypatch.setenv('TWINKLE_NODE_RANKS', '0,2') + monkeypatch.setattr(dist, 'get_rank', lambda: 2) + monkeypatch.setattr(dist, 'get_world_size', lambda: 3) + + assert _get_local_rank_info() == (2, 3, 0, [0, 2]) + + +def test_native_fsdp_keeps_torchrun_local_rank_fallback(monkeypatch): + _clear_topology_env(monkeypatch) + monkeypatch.setenv('LOCAL_RANK', '1') + monkeypatch.setenv('LOCAL_WORLD_SIZE', '2') + monkeypatch.setattr(dist, 'get_rank', lambda: 1) + monkeypatch.setattr(dist, 'get_world_size', lambda: 2) + + assert _get_local_rank_info() == (1, 2, 0, [0, 1]) From c3bce6b200cc126eab376201dfde7e343bdcd894 Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Thu, 13 Aug 2026 18:27:57 +0800 Subject: [PATCH 05/28] wip --- .../server/transformer/run_dsv4_0731_npu.sh | 67 ++++++++++ .../server_config_dsv4_0731_npu.yaml | 116 ++++++++++++++++++ .../transformer/server_dsv4_0731_npu.py | 12 ++ .../twinkle/run_dsv4_0731_npu_client.sh | 26 ++++ 4 files changed, 221 insertions(+) create mode 100755 cookbook/client/server/transformer/run_dsv4_0731_npu.sh create mode 100644 cookbook/client/server/transformer/server_config_dsv4_0731_npu.yaml create mode 100644 cookbook/client/server/transformer/server_dsv4_0731_npu.py create mode 100755 cookbook/client/twinkle/run_dsv4_0731_npu_client.sh diff --git a/cookbook/client/server/transformer/run_dsv4_0731_npu.sh b/cookbook/client/server/transformer/run_dsv4_0731_npu.sh new file mode 100755 index 000000000..dd2eea223 --- /dev/null +++ b/cookbook/client/server/transformer/run_dsv4_0731_npu.sh @@ -0,0 +1,67 @@ +#!/usr/bin/env bash +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +PROJECT_DIR="$(cd -- "$SCRIPT_DIR/../../../.." && pwd)" +cd "$PROJECT_DIR" + +if [[ -f /usr/local/Ascend/ascend-toolkit/set_env.sh ]]; then + # shellcheck disable=SC1091 + source /usr/local/Ascend/ascend-toolkit/set_env.sh +fi + +export PYTHONPATH="$PROJECT_DIR/src${PYTHONPATH:+:$PYTHONPATH}" +export ASCEND_RT_VISIBLE_DEVICES="${ASCEND_RT_VISIBLE_DEVICES:-0,1}" +export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1 +export TWINKLE_TRUST_REMOTE_CODE=1 +export TWINKLE_FAIL_FAST=1 +export TOKENIZERS_PARALLELISM=true +export HCCL_CONNECT_TIMEOUT="${HCCL_CONNECT_TIMEOUT:-7200}" +export HCCL_EXEC_TIMEOUT="${HCCL_EXEC_TIMEOUT:-0}" + +MODEL_PATH=/nas/disk1/random-deepseek-v4-4b +CONFIG_PATH="$SCRIPT_DIR/server_config_dsv4_0731_npu.yaml" + +test -f "$MODEL_PATH/config.json" +test -f "$MODEL_PATH/tokenizer.json" + +python3 - <<'PY' +import os +import torch +import torch_npu # noqa: F401 + +visible = [item.strip() for item in os.environ['ASCEND_RT_VISIBLE_DEVICES'].split(',') if item.strip()] +if len(visible) != 2: + raise SystemExit(f'需要两张 NPU，当前 ASCEND_RT_VISIBLE_DEVICES={visible}') +if not torch.npu.is_available(): + raise SystemExit('torch.npu.is_available() 为 False') +if torch.npu.device_count() < 2: + raise SystemExit(f'当前进程只能看到 {torch.npu.device_count()} 张 NPU') +print(f'NPU 检查通过：visible={visible}, device_count={torch.npu.device_count()}') +PY + +if ! ray status >/dev/null 2>&1; then + ray start \ + --head \ + --num-cpus="${TWINKLE_RAY_CPUS:-8}" \ + --resources='{"NPU": 2}' \ + --disable-usage-stats \ + --include-dashboard=false +fi + +python3 - <<'PY' +import ray + +ray.init(address='auto', logging_level='ERROR') +npu_count = float(ray.cluster_resources().get('NPU', 0)) +ray.shutdown() +if npu_count < 2: + raise SystemExit( + f'当前 Ray 集群只有 {npu_count:g} 个 NPU 资源。' + '请确认没有其他任务占用后执行 ray stop --force，再重新运行本脚本。' + ) +print(f'Ray NPU 资源检查通过：NPU={npu_count:g}') +PY + +python3 -m twinkle.server check-config --config "$CONFIG_PATH" +exec python3 "$SCRIPT_DIR/server_dsv4_0731_npu.py" diff --git a/cookbook/client/server/transformer/server_config_dsv4_0731_npu.yaml b/cookbook/client/server/transformer/server_config_dsv4_0731_npu.yaml new file mode 100644 index 000000000..6d00f66eb --- /dev/null +++ b/cookbook/client/server/transformer/server_config_dsv4_0731_npu.yaml @@ -0,0 +1,116 @@ +# DeepSeek-V4-Flash-0731 client/server Multi-LoRA SFT on one 2-NPU node. +# Both NPUs are assigned to the Transformers training model. A sampler is +# intentionally omitted because it would compete with training for devices. + +proxy_location: EveryNode + +http_options: + host: 0.0.0.0 + port: 8000 + +telemetry: + enabled: false + otlp_endpoint: http://localhost:4317 + +persistence: + mode: file + file_path: /tmp/twinkle_dsv4_0731_npu_state.json + +applications: + - name: server + route_prefix: /api/v1 + import_path: server + args: + server_config: + per_token_model_limit: 2 + supported_models: + - deepseek-v4-0731-local + deployments: + - name: TinkerCompatServer + max_ongoing_requests: 50 + autoscaling_config: + min_replicas: 1 + max_replicas: 1 + target_ongoing_requests: 32 + ray_actor_options: + num_cpus: 0.1 + runtime_env: + env_vars: + TWINKLE_FAIL_FAST: "1" + + - name: models-deepseek-v4-0731-local + route_prefix: /api/v1/model/deepseek-v4-0731-local + import_path: model + args: + backend: transformers + model_id: /nas/disk1/random-deepseek-v4-4b + nproc_per_node: 2 + max_loras: 2 + max_r: 8 + max_length: 2048 + mixed_precision: bf16 + strategy: native_fsdp + memory_efficient_init: true + target_modules: all-linear + fsdp_config: + reshard_after_forward: true + expert_parallel: + enabled: true + ep_size: 2 + router_dtype: fp32 + keep_router_logits: false + device_group: + name: model + ranks: 2 + device_type: npu + device_mesh: + device_type: npu + fsdp_size: 2 + dp_size: 1 + ep_size: 2 + queue_config: + rps_limit: 20 + tps_limit: 100000 + max_input_tokens: 4096 + queue_timeout: 3600 + execution_timeout: 3600 + adapter_config: + adapter_timeout: 3600 + adapter_max_lifetime: 86400 + deployments: + - name: ModelManagement + autoscaling_config: + min_replicas: 1 + max_replicas: 1 + target_ongoing_requests: 8 + ray_actor_options: + num_cpus: 0.1 + runtime_env: + env_vars: + TWINKLE_TRUST_REMOTE_CODE: "1" + TWINKLE_FAIL_FAST: "1" + + - name: processor + route_prefix: /api/v1/processor + import_path: processor + args: + ncpu_proc_per_node: 2 + device_group: + name: processor + ranks: 2 + device_type: CPU + device_mesh: + device_type: CPU + dp_size: 2 + deployments: + - name: ProcessorManagement + autoscaling_config: + min_replicas: 1 + max_replicas: 1 + target_ongoing_requests: 32 + ray_actor_options: + num_cpus: 0.1 + runtime_env: + env_vars: + TWINKLE_TRUST_REMOTE_CODE: "1" + TWINKLE_FAIL_FAST: "1" diff --git a/cookbook/client/server/transformer/server_dsv4_0731_npu.py b/cookbook/client/server/transformer/server_dsv4_0731_npu.py new file mode 100644 index 000000000..7b19f24c3 --- /dev/null +++ b/cookbook/client/server/transformer/server_dsv4_0731_npu.py @@ -0,0 +1,12 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Launch the two-NPU DeepSeek-V4-0731 Multi-LoRA server.""" +import os + +os.environ.setdefault('TWINKLE_TRUST_REMOTE_CODE', '1') + +from twinkle.server import launch_server # noqa: E402 + +file_dir = os.path.abspath(os.path.dirname(__file__)) +config_path = os.path.join(file_dir, 'server_config_dsv4_0731_npu.yaml') + +launch_server(config_path=config_path) diff --git a/cookbook/client/twinkle/run_dsv4_0731_npu_client.sh b/cookbook/client/twinkle/run_dsv4_0731_npu_client.sh new file mode 100755 index 000000000..60068b705 --- /dev/null +++ b/cookbook/client/twinkle/run_dsv4_0731_npu_client.sh @@ -0,0 +1,26 @@ +#!/usr/bin/env bash +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +PROJECT_DIR="$(cd -- "$SCRIPT_DIR/../../.." && pwd)" +cd "$PROJECT_DIR" + +export PYTHONPATH="$PROJECT_DIR/src${PYTHONPATH:+:$PYTHONPATH}" +export TWINKLE_SERVER_URL="${TWINKLE_SERVER_URL:-http://127.0.0.1:8000}" +export TWINKLE_SERVER_TOKEN="${TWINKLE_SERVER_TOKEN:-EMPTY_TOKEN}" +export TWINKLE_MODEL_ID="${TWINKLE_MODEL_ID:-deepseek-v4-0731-local}" +export DSV4_MODEL_ID="${DSV4_MODEL_ID:-/nas/disk1/random-deepseek-v4-4b}" +export DATASET_ID="${DATASET_ID:-/model/ljl/dataset/self-cognition.jsonl}" +export ADAPTER_NAMES="${ADAPTER_NAMES:-tenant_a,tenant_b}" +export BATCH_SIZE="${BATCH_SIZE:-2}" +export GRAD_ACCUM_STEPS="${GRAD_ACCUM_STEPS:-4}" +export MAX_STEPS="${MAX_STEPS:-10}" +export MAX_LENGTH="${MAX_LENGTH:-2048}" +export LORA_R="${LORA_R:-8}" +export LORA_ALPHA="${LORA_ALPHA:-32}" +export LR="${LR:-1e-4}" +export OUTPUT_DIR="${OUTPUT_DIR:-/tmp/twinkle_dsv4_0731_npu_multi_lora}" + +test -f "$DATASET_ID" + +exec python3 "$SCRIPT_DIR/dsv4_multi_lora_sft.py" From 8c0aa5b2f71cd36ac8c72c5d51189829885880aa Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Fri, 14 Aug 2026 10:36:34 +0800 Subject: [PATCH 06/28] wip --- .../model/megatron/multi_lora_megatron.py | 3 +- src/twinkle/model/multi_lora.py | 20 +++-- .../model/multi_lora_target_parameters.py | 6 +- .../transformers/multi_lora_transformers.py | 3 +- .../model/test_multi_lora_dtensor_release.py | 77 +++++++++++++++++++ .../test_multi_lora_target_parameters.py | 2 +- 6 files changed, 99 insertions(+), 12 deletions(-) create mode 100644 tests/model/test_multi_lora_dtensor_release.py diff --git a/src/twinkle/model/megatron/multi_lora_megatron.py b/src/twinkle/model/megatron/multi_lora_megatron.py index ebda91501..6095d2cef 100644 --- a/src/twinkle/model/megatron/multi_lora_megatron.py +++ b/src/twinkle/model/megatron/multi_lora_megatron.py @@ -446,6 +446,5 @@ def add_metric(self, metric_cls: Union[Metric, str], is_training: Optional[bool] @remote_function() def remove_adapter(self, adapter_name: str): - if adapter_name in self.optimizer_group: - self.optimizer_group.pop(adapter_name) self.multi_adapter.release_lora(adapter_name) + self.optimizer_group.pop(adapter_name, None) diff --git a/src/twinkle/model/multi_lora.py b/src/twinkle/model/multi_lora.py index 0a66b81ca..59a3914d7 100644 --- a/src/twinkle/model/multi_lora.py +++ b/src/twinkle/model/multi_lora.py @@ -66,7 +66,14 @@ def _is_distributed_param(parameter): def _is_target_parameter_lora_name(name: str) -> bool: return '._twinkle_lora_' in name + @torch.no_grad() def _write_param_tensor(self, parameter, value): + """Copy a value into a regular parameter or its local DTensor shard. + + ``DTensor.to_local()`` can return a view produced by a custom autograd + Function. In-place writes to that view are forbidden while grad mode is + enabled, even for lifecycle operations such as adapter reset/load. + """ if value is None: return value = value.detach().to(dtype=parameter.dtype) @@ -238,13 +245,16 @@ def acquire_lora(self, tenant_adapter_name: str, config: LoraConfig) -> str: def release_lora(self, tenant_adapter_name: str) -> Optional[str]: try: _lora = self.find_lora_by_tenant(tenant_adapter_name) - _lora.tenant_config = None - _lora.tenant_adapter_name = None - self._load_initial_weights(_lora.adapter_name) - self.target_parameter_manager.release(tenant_adapter_name) - logger.info(f'Lora count: {len(self.loras)}, available lora: {self._count_available_loras()}') except ValueError: return + # Restore every backing slot before publishing it as available. If a + # DTensor reset fails, retain the tenant mapping so cleanup can retry + # safely instead of exposing a partially reset LoRA slot. + self._load_initial_weights(_lora.adapter_name) + self.target_parameter_manager.release(tenant_adapter_name) + _lora.tenant_config = None + _lora.tenant_adapter_name = None + logger.info(f'Lora count: {len(self.loras)}, available lora: {self._count_available_loras()}') def has_lora(self, adapter_name: str) -> bool: return len([_lora for _lora in self.loras if _lora.tenant_adapter_name == adapter_name]) > 0 diff --git a/src/twinkle/model/multi_lora_target_parameters.py b/src/twinkle/model/multi_lora_target_parameters.py index b8af6c4d2..73ee35488 100644 --- a/src/twinkle/model/multi_lora_target_parameters.py +++ b/src/twinkle/model/multi_lora_target_parameters.py @@ -133,6 +133,7 @@ def _read_parameter(parameter: nn.Parameter) -> torch.Tensor: return parameter @staticmethod + @torch.no_grad() def _write_parameter(parameter: nn.Parameter, value: torch.Tensor) -> None: if hasattr(parameter, 'to_local') and hasattr(parameter, 'device_mesh'): local_parameter = parameter.to_local() @@ -401,12 +402,13 @@ def acquire(self, tenant_adapter_name: str, slot_name: str, config: LoraConfig) wrapper.configure_slot(slot_name, config) def release(self, tenant_adapter_name: str) -> None: - slot_name = self.tenant_to_slot.pop(tenant_adapter_name, None) - self.tenant_configs.pop(tenant_adapter_name, None) + slot_name = self.tenant_to_slot.get(tenant_adapter_name) if slot_name is None: return for wrapper in self.wrappers: wrapper.reset_slot(slot_name) + self.tenant_to_slot.pop(tenant_adapter_name, None) + self.tenant_configs.pop(tenant_adapter_name, None) def save_initial_weights(self) -> None: for wrapper in self.wrappers: diff --git a/src/twinkle/model/transformers/multi_lora_transformers.py b/src/twinkle/model/transformers/multi_lora_transformers.py index 0f95d937d..dc941bc9a 100644 --- a/src/twinkle/model/transformers/multi_lora_transformers.py +++ b/src/twinkle/model/transformers/multi_lora_transformers.py @@ -349,9 +349,8 @@ def calculate_metric(self, is_training, **kwargs): @remote_function() def remove_adapter(self, adapter_name: str): self._lazy_wrap_model() - if adapter_name in self.optimizer_group: - self.optimizer_group.pop(adapter_name) self.multi_adapter.release_lora(adapter_name) + self.optimizer_group.pop(adapter_name, None) def _get_nb_trainable_parameters(self, adapter_name, model): with self.multi_adapter.adapter(adapter_name): diff --git a/tests/model/test_multi_lora_dtensor_release.py b/tests/model/test_multi_lora_dtensor_release.py new file mode 100644 index 000000000..b225810ae --- /dev/null +++ b/tests/model/test_multi_lora_dtensor_release.py @@ -0,0 +1,77 @@ +import pytest +import torch +from peft import LoraConfig +from torch import nn + +from twinkle.model.multi_lora import LoraTenant, MultiLora +from twinkle.model.multi_lora_target_parameters import TargetParameterLoraWrapper + + +class _AutogradView(torch.autograd.Function): + + @staticmethod + def forward(ctx, tensor): + return tensor.view_as(tensor) + + @staticmethod + def backward(ctx, grad): + return grad + + +class _FakeDistributedParameter: + + def __init__(self, local_parameter, global_shape): + self._local_parameter = local_parameter + self.shape = global_shape + self.dtype = local_parameter.dtype + self.device = local_parameter.device + self.device_mesh = object() + self.placements = () + + def to_local(self): + return _AutogradView.apply(self._local_parameter) + + +def _make_lora_config(rank): + return LoraConfig(r=rank, lora_alpha=rank * 2, target_modules=['linear']) + + +def test_multilora_writes_dtensor_local_autograd_view_without_tracking_gradients(): + local_parameter = nn.Parameter(torch.ones(2)) + distributed_parameter = _FakeDistributedParameter(local_parameter, global_shape=(4,)) + + MultiLora()._write_param_tensor(distributed_parameter, torch.zeros(2)) + + assert torch.count_nonzero(local_parameter) == 0 + + +def test_target_parameter_lora_writes_dtensor_local_autograd_view_without_tracking_gradients(): + local_parameter = nn.Parameter(torch.ones(2)) + distributed_parameter = _FakeDistributedParameter(local_parameter, global_shape=(4,)) + + TargetParameterLoraWrapper._write_parameter(distributed_parameter, torch.zeros(2)) + + assert torch.count_nonzero(local_parameter) == 0 + + +def test_multilora_release_keeps_tenant_when_slot_reset_fails(monkeypatch): + multi_lora = MultiLora(max_loras=1, max_r=4) + tenant = LoraTenant( + index=0, + adapter_name='lora_0', + config=_make_lora_config(4), + tenant_adapter_name='adapter_a', + tenant_config=_make_lora_config(2), + ) + multi_lora.loras = [tenant] + + def fail_reset(_adapter_name): + raise RuntimeError('reset failed') + + monkeypatch.setattr(multi_lora, '_load_initial_weights', fail_reset) + + with pytest.raises(RuntimeError, match='reset failed'): + multi_lora.release_lora('adapter_a') + + assert tenant.tenant_adapter_name == 'adapter_a' + assert tenant.tenant_config is not None diff --git a/tests/model/test_multi_lora_target_parameters.py b/tests/model/test_multi_lora_target_parameters.py index b28ef6b36..fb47e39bd 100644 --- a/tests/model/test_multi_lora_target_parameters.py +++ b/tests/model/test_multi_lora_target_parameters.py @@ -266,4 +266,4 @@ def test_multilora_transformers_installs_target_parameters_once(): assert test_target_parameter_multi_lora_updates_only_active_adapter() == True assert test_multilora_releases_target_parameter_slot_to_initial_weights() == True assert test_multilora_state_dict_round_trips_target_parameters() == True - assert test_multilora_transformers_installs_target_parameters_once() == True \ No newline at end of file + assert test_multilora_transformers_installs_target_parameters_once() == True From 9ad1524a9dfd7322f5b2e3e4da1728d7ade61ed4 Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Fri, 14 Aug 2026 11:52:53 +0800 Subject: [PATCH 07/28] wip --- src/twinkle/model/multi_lora.py | 25 ++++++++++++++-- .../model/test_multi_lora_dtensor_release.py | 29 +++++++++++++++++++ 2 files changed, 52 insertions(+), 2 deletions(-) diff --git a/src/twinkle/model/multi_lora.py b/src/twinkle/model/multi_lora.py index 59a3914d7..529dbcd65 100644 --- a/src/twinkle/model/multi_lora.py +++ b/src/twinkle/model/multi_lora.py @@ -132,6 +132,10 @@ def _copy_rank_tensor(name: str, target, value): def _count_available_loras(self): return len([_lora for _lora in self.loras if _lora.tenant_adapter_name is None]) + def _lora_slot_assignments(self): + """Return a compact snapshot used to diagnose per-rank slot state.""" + return [(lora.adapter_name, lora.tenant_adapter_name) for lora in self.loras] + def reset_adapter_status(self): """Force lora_0 require_grad, disable others""" if isinstance(self.module, list): @@ -239,13 +243,24 @@ def acquire_lora(self, tenant_adapter_name: str, config: LoraConfig) -> str: slot_name=_available_lora.adapter_name, config=config, ) - logger.info(f'Lora count: {len(self.loras)}, available lora: {self._count_available_loras()}') + logger.info( + 'LoRA acquired: tenant=%s, slot=%s, available_lora=%s', + tenant_adapter_name, + _available_lora.adapter_name, + self._count_available_loras(), + ) return _available_lora.adapter_name def release_lora(self, tenant_adapter_name: str) -> Optional[str]: try: _lora = self.find_lora_by_tenant(tenant_adapter_name) except ValueError: + logger.warning( + 'LoRA release skipped: tenant=%s was not found, assignments=%s', + tenant_adapter_name, + self._lora_slot_assignments(), + ranks='all', + ) return # Restore every backing slot before publishing it as available. If a # DTensor reset fails, retain the tenant mapping so cleanup can retry @@ -254,7 +269,13 @@ def release_lora(self, tenant_adapter_name: str) -> Optional[str]: self.target_parameter_manager.release(tenant_adapter_name) _lora.tenant_config = None _lora.tenant_adapter_name = None - logger.info(f'Lora count: {len(self.loras)}, available lora: {self._count_available_loras()}') + logger.info( + 'LoRA released: tenant=%s, slot=%s, available_lora=%s', + tenant_adapter_name, + _lora.adapter_name, + self._count_available_loras(), + ) + return _lora.adapter_name def has_lora(self, adapter_name: str) -> bool: return len([_lora for _lora in self.loras if _lora.tenant_adapter_name == adapter_name]) > 0 diff --git a/tests/model/test_multi_lora_dtensor_release.py b/tests/model/test_multi_lora_dtensor_release.py index b225810ae..9794726fa 100644 --- a/tests/model/test_multi_lora_dtensor_release.py +++ b/tests/model/test_multi_lora_dtensor_release.py @@ -75,3 +75,32 @@ def fail_reset(_adapter_name): assert tenant.tenant_adapter_name == 'adapter_a' assert tenant.tenant_config is not None + + +def test_multilora_release_reports_each_released_slot(monkeypatch): + multi_lora = MultiLora(max_loras=2, max_r=4) + slot_config = _make_lora_config(4) + tenant_config = _make_lora_config(2) + multi_lora.loras = [ + LoraTenant( + index=0, + adapter_name='lora_0', + config=slot_config, + tenant_adapter_name='adapter_a', + tenant_config=tenant_config, + ), + LoraTenant( + index=1, + adapter_name='lora_1', + config=slot_config, + tenant_adapter_name='adapter_b', + tenant_config=tenant_config, + ), + ] + monkeypatch.setattr(multi_lora, '_load_initial_weights', lambda _adapter_name: None) + monkeypatch.setattr(multi_lora.target_parameter_manager, 'release', lambda _tenant_name: None) + + assert multi_lora.release_lora('adapter_a') == 'lora_0' + assert multi_lora._count_available_loras() == 1 + assert multi_lora.release_lora('adapter_b') == 'lora_1' + assert multi_lora._count_available_loras() == 2 From ad7b9f4aefc5c63a315ba24c37fc9744414b82c1 Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Fri, 14 Aug 2026 15:01:07 +0800 Subject: [PATCH 08/28] multinode script --- .../run_dsv4_0731_npu_multinode.sh | 154 ++++++++++++++++++ ...server_config_dsv4_0731_npu_multinode.yaml | 126 ++++++++++++++ .../run_dsv4_0731_npu_multinode_client.sh | 37 +++++ 3 files changed, 317 insertions(+) create mode 100755 cookbook/client/server/transformer/run_dsv4_0731_npu_multinode.sh create mode 100644 cookbook/client/server/transformer/server_config_dsv4_0731_npu_multinode.yaml create mode 100755 cookbook/client/twinkle/run_dsv4_0731_npu_multinode_client.sh diff --git a/cookbook/client/server/transformer/run_dsv4_0731_npu_multinode.sh b/cookbook/client/server/transformer/run_dsv4_0731_npu_multinode.sh new file mode 100755 index 000000000..55a5aaba1 --- /dev/null +++ b/cookbook/client/server/transformer/run_dsv4_0731_npu_multinode.sh @@ -0,0 +1,154 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Start a two-node Ascend A3 Ray cluster and launch the 32-NPU Twinkle server. +# +# Head node (starts Ray, waits for the worker, then launches Twinkle Server): +# HEAD_IP=10.0.0.10 NODE_IP=10.0.0.10 NETWORK_IFACE=eth0 \ +# bash run_dsv4_0731_npu_multinode.sh head +# +# Worker node (joins Ray and exits, leaving the Ray daemon running): +# HEAD_IP=10.0.0.10 NODE_IP=10.0.0.11 NETWORK_IFACE=eth0 \ +# bash run_dsv4_0731_npu_multinode.sh worker +# +# Set RESET_RAY=1 when intentionally replacing an existing local Ray runtime. + +ROLE="${1:-}" +if [[ "$ROLE" != "head" && "$ROLE" != "worker" ]]; then + echo "Usage: HEAD_IP= NODE_IP= NETWORK_IFACE= $0 {head|worker}" >&2 + exit 2 +fi + +: "${HEAD_IP:?Set HEAD_IP to the Ray head node IP}" +: "${NODE_IP:?Set NODE_IP to this node IP}" + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +PROJECT_DIR="$(cd -- "$SCRIPT_DIR/../../../.." && pwd)" +CONFIG_PATH="$SCRIPT_DIR/server_config_dsv4_0731_npu_multinode.yaml" +MODEL_ID="hf://deepseek-ai/DeepSeek-V4-Flash-0731" +DATASET_PATH="${DATASET_ID:-/model/ljl/dataset/self-cognition.jsonl}" + +NPU_PER_NODE=16 +NNODES=2 +TOTAL_NPUS=$((NPU_PER_NODE * NNODES)) +RAY_PORT="${RAY_PORT:-6379}" +RAY_CPUS_PER_NODE="${TWINKLE_RAY_CPUS:-$(nproc)}" +CLUSTER_WAIT_SECONDS="${CLUSTER_WAIT_SECONDS:-1800}" +NETWORK_IFACE="${NETWORK_IFACE:-eth0}" + +if [[ -f /usr/local/Ascend/ascend-toolkit/set_env.sh ]]; then + # shellcheck disable=SC1091 + source /usr/local/Ascend/ascend-toolkit/set_env.sh +fi + +cd "$PROJECT_DIR" +export PYTHONPATH="$PROJECT_DIR/src${PYTHONPATH:+:$PYTHONPATH}" +export ASCEND_RT_VISIBLE_DEVICES="${ASCEND_RT_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15}" +export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1 +export TWINKLE_TRUST_REMOTE_CODE=1 +export TWINKLE_FAIL_FAST=1 +export TOKENIZERS_PARALLELISM=true +export GLOO_SOCKET_IFNAME="$NETWORK_IFACE" +export HCCL_SOCKET_IFNAME="$NETWORK_IFACE" +export HCCL_CONNECT_TIMEOUT="${HCCL_CONNECT_TIMEOUT:-7200}" +export HCCL_EXEC_TIMEOUT="${HCCL_EXEC_TIMEOUT:-0}" +export HCCL_IF_BASE_PORT="${HCCL_IF_BASE_PORT:-20000}" + +if [[ "$ROLE" == "head" ]]; then + export NODE_RANK=0 +else + export NODE_RANK=1 +fi + +IFS=',' read -r -a VISIBLE_NPUS <<< "$ASCEND_RT_VISIBLE_DEVICES" +if [[ "${#VISIBLE_NPUS[@]}" -ne "$NPU_PER_NODE" ]]; then + echo "Expected $NPU_PER_NODE visible NPUs, got ASCEND_RT_VISIBLE_DEVICES=$ASCEND_RT_VISIBLE_DEVICES" >&2 + exit 1 +fi + +test -f "$CONFIG_PATH" +test -f "$DATASET_PATH" + +if [[ "$MODEL_ID" == hf://* || "$MODEL_ID" == ms://* ]]; then + echo "Model will be downloaded through the configured Hub backend: $MODEL_ID" +elif [[ -d "$MODEL_ID" ]]; then + test -f "$MODEL_ID/config.json" + test -f "$MODEL_ID/tokenizer.json" +else + echo "Invalid MODEL_ID: use an existing local directory or an explicit hf:// or ms:// ID." >&2 + exit 1 +fi + +python3 - "$NPU_PER_NODE" <<'PY' +import sys +import torch +import torch_npu # noqa: F401 + +expected = int(sys.argv[1]) +if not torch.npu.is_available(): + raise SystemExit('torch.npu.is_available() is False') +actual = torch.npu.device_count() +if actual < expected: + raise SystemExit(f'Expected at least {expected} visible NPUs, got {actual}') +print(f'Ascend check passed: visible NPU count={actual}') +PY + +if [[ "${RESET_RAY:-0}" == "1" ]]; then + ray stop --force || true +fi + +if ray status >/dev/null 2>&1; then + echo "Ray is already running on this node; set RESET_RAY=1 to replace it." +elif [[ "$ROLE" == "head" ]]; then + ray start \ + --head \ + --node-ip-address="$HEAD_IP" \ + --port="$RAY_PORT" \ + --num-cpus="$RAY_CPUS_PER_NODE" \ + --resources="{\"NPU\": $NPU_PER_NODE}" \ + --disable-usage-stats \ + --include-dashboard=false +else + ray start \ + --address="$HEAD_IP:$RAY_PORT" \ + --node-ip-address="$NODE_IP" \ + --num-cpus="$RAY_CPUS_PER_NODE" \ + --resources="{\"NPU\": $NPU_PER_NODE}" \ + --disable-usage-stats +fi + +if [[ "$ROLE" == "worker" ]]; then + echo "Worker joined Ray at $HEAD_IP:$RAY_PORT with $NPU_PER_NODE NPU resources." + exit 0 +fi + +python3 - "$TOTAL_NPUS" "$NNODES" "$NPU_PER_NODE" "$CLUSTER_WAIT_SECONDS" <<'PY' +import sys +import time +import ray + +expected_total = int(sys.argv[1]) +expected_nodes = int(sys.argv[2]) +expected_per_node = int(sys.argv[3]) +timeout = int(sys.argv[4]) +deadline = time.monotonic() + timeout + +ray.init(address='auto', logging_level='ERROR') +try: + while True: + alive_nodes = [node for node in ray.nodes() if node.get('Alive', True)] + npu_nodes = [node for node in alive_nodes if float(node.get('Resources', {}).get('NPU', 0)) >= expected_per_node] + total = int(ray.cluster_resources().get('NPU', 0)) + print(f'Waiting for cluster: NPU={total}/{expected_total}, NPU nodes={len(npu_nodes)}/{expected_nodes}') + if total >= expected_total and len(npu_nodes) >= expected_nodes: + break + if time.monotonic() >= deadline: + raise SystemExit('Timed out waiting for the two-node 32-NPU Ray cluster') + time.sleep(5) +finally: + ray.shutdown() +PY + +python3 -m twinkle.server check-config --config "$CONFIG_PATH" +echo "Launching Twinkle Server at http://$HEAD_IP:8000" +exec python3 -m twinkle.server launch --config "$CONFIG_PATH" diff --git a/cookbook/client/server/transformer/server_config_dsv4_0731_npu_multinode.yaml b/cookbook/client/server/transformer/server_config_dsv4_0731_npu_multinode.yaml new file mode 100644 index 000000000..955f98bf6 --- /dev/null +++ b/cookbook/client/server/transformer/server_config_dsv4_0731_npu_multinode.yaml @@ -0,0 +1,126 @@ +# DeepSeek-V4-Flash-0731 Multi-LoRA SFT on two Ascend A3 nodes. +# Each node contributes 16 NPUs, for a 32-rank Native FSDP2 + EP deployment. +# +# Required on both nodes: +# - the same Twinkle commit and Python/CANN/torch_npu environment +# - enough local Hub cache space on both nodes for the downloaded checkpoint +# - OUTPUT_DIR used by the client mounted as a shared filesystem + +proxy_location: EveryNode + +http_options: + host: 0.0.0.0 + port: 8000 + +telemetry: + enabled: false + otlp_endpoint: http://localhost:4317 + +# "memory" is backed by one detached Ray actor and is shared by all Serve +# processes on both nodes. It does not survive a full Ray-cluster shutdown. +persistence: + mode: memory + key_prefix: dsv4-0731-a3-32npu + +applications: + - name: server + route_prefix: /api/v1 + import_path: server + args: + server_config: + per_token_model_limit: 2 + supported_models: + - deepseek-v4-0731-local + deployments: + - name: TinkerCompatServer + max_ongoing_requests: 50 + autoscaling_config: + min_replicas: 1 + max_replicas: 1 + target_ongoing_requests: 32 + ray_actor_options: + num_cpus: 0.1 + runtime_env: + env_vars: + TWINKLE_FAIL_FAST: "1" + + - name: models-deepseek-v4-0731-local + route_prefix: /api/v1/model/deepseek-v4-0731-local + import_path: model + args: + backend: transformers + # Twinkle treats an unprefixed ID as ModelScope. Keep the hf:// prefix + # here to select Hugging Face explicitly. + model_id: hf://deepseek-ai/DeepSeek-V4-Flash-0731 + nproc_per_node: 16 + max_loras: 2 + max_r: 8 + max_length: 2048 + mixed_precision: bf16 + strategy: native_fsdp + memory_efficient_init: true + target_modules: all-linear + fsdp_config: + reshard_after_forward: true + expert_parallel: + enabled: true + ep_size: 32 + router_dtype: fp32 + keep_router_logits: false + device_group: + name: model + ranks: 32 + device_type: npu + device_mesh: + device_type: npu + fsdp_size: 32 + dp_size: 1 + ep_size: 32 + queue_config: + rps_limit: 20 + tps_limit: 100000 + max_input_tokens: 4096 + queue_timeout: 7200 + execution_timeout: 7200 + adapter_config: + adapter_timeout: 60 + adapter_max_lifetime: 86400 + deployments: + - name: ModelManagement + autoscaling_config: + min_replicas: 1 + max_replicas: 1 + target_ongoing_requests: 8 + ray_actor_options: + num_cpus: 0.1 + runtime_env: + env_vars: + TWINKLE_TRUST_REMOTE_CODE: "1" + TWINKLE_FAIL_FAST: "1" + + # Dataset construction and tokenization run on CPU. The dataset and model + # encoding files are resolved on the node selected for this deployment. + - name: processor + route_prefix: /api/v1/processor + import_path: processor + args: + ncpu_proc_per_node: 2 + device_group: + name: processor + ranks: 2 + device_type: CPU + device_mesh: + device_type: CPU + dp_size: 2 + deployments: + - name: ProcessorManagement + autoscaling_config: + min_replicas: 1 + max_replicas: 1 + target_ongoing_requests: 32 + ray_actor_options: + num_cpus: 0.1 + runtime_env: + env_vars: + TWINKLE_TRUST_REMOTE_CODE: "1" + TWINKLE_FAIL_FAST: "1" diff --git a/cookbook/client/twinkle/run_dsv4_0731_npu_multinode_client.sh b/cookbook/client/twinkle/run_dsv4_0731_npu_multinode_client.sh new file mode 100755 index 000000000..1e0ff2006 --- /dev/null +++ b/cookbook/client/twinkle/run_dsv4_0731_npu_multinode_client.sh @@ -0,0 +1,37 @@ +#!/usr/bin/env bash +set -euo pipefail + +# The client is an HTTP orchestrator. Run exactly one client process; it does +# not join Ray and does not need NPU devices. +# +# Example: +# HEAD_IP=10.0.0.10 \ +# OUTPUT_DIR=/shared/twinkle_output/dsv4-0731-a3-32npu \ +# bash run_dsv4_0731_npu_multinode_client.sh + +: "${HEAD_IP:?Set HEAD_IP to the Twinkle/Ray head node IP}" +: "${OUTPUT_DIR:?Set OUTPUT_DIR to a shared absolute path mounted on both server nodes}" + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +PROJECT_DIR="$(cd -- "$SCRIPT_DIR/../../.." && pwd)" +cd "$PROJECT_DIR" + +export PYTHONPATH="$PROJECT_DIR/src${PYTHONPATH:+:$PYTHONPATH}" +export TWINKLE_SERVER_URL="${TWINKLE_SERVER_URL:-http://$HEAD_IP:8000}" +export TWINKLE_SERVER_TOKEN="${TWINKLE_SERVER_TOKEN:-EMPTY_TOKEN}" +export TWINKLE_MODEL_ID="${TWINKLE_MODEL_ID:-deepseek-v4-0731-local}" +# The template/encoding side must resolve the same checkpoint as the server. +export DSV4_MODEL_ID="${DSV4_MODEL_ID:-hf://deepseek-ai/DeepSeek-V4-Flash-0731}" +export DATASET_ID="${DATASET_ID:-/model/ljl/dataset/self-cognition.jsonl}" +export OUTPUT_DIR + +# forward_backward dispatches one slice to each of the 32 FSDP ranks. +export BATCH_SIZE="${BATCH_SIZE:-32}" +export GRAD_ACCUM_STEPS="${GRAD_ACCUM_STEPS:-4}" +export MAX_STEPS="${MAX_STEPS:-10}" +export MAX_LENGTH="${MAX_LENGTH:-2048}" +export LORA_R="${LORA_R:-8}" +export LORA_ALPHA="${LORA_ALPHA:-32}" + +mkdir -p "$OUTPUT_DIR" +exec python3 cookbook/client/twinkle/dsv4_multi_lora_sft.py From 71d3007c20e3a31d0bd7d293a9f07c24ea4ef2d1 Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Fri, 14 Aug 2026 17:20:08 +0800 Subject: [PATCH 09/28] wip --- cookbook/client/server/transformer/read_me.md | 47 +++++ .../run_dsv4_0731_npu_2node_2npu.sh | 161 ++++++++++++++++++ ...erver_config_dsv4_0731_npu_2node_2npu.yaml | 121 +++++++++++++ .../run_dsv4_0731_npu_2node_2npu_client.sh | 33 ++++ 4 files changed, 362 insertions(+) create mode 100644 cookbook/client/server/transformer/read_me.md create mode 100755 cookbook/client/server/transformer/run_dsv4_0731_npu_2node_2npu.sh create mode 100644 cookbook/client/server/transformer/server_config_dsv4_0731_npu_2node_2npu.yaml create mode 100755 cookbook/client/twinkle/run_dsv4_0731_npu_2node_2npu_client.sh diff --git a/cookbook/client/server/transformer/read_me.md b/cookbook/client/server/transformer/read_me.md new file mode 100644 index 000000000..57d239056 --- /dev/null +++ b/cookbook/client/server/transformer/read_me.md @@ -0,0 +1,47 @@ +已新增独立的“两台机器、每台 2 张昇腾 A3”配置，原来的 16 卡文件没有修改。 + +文件： + +- [4 卡服务端配置](/Users/linjiajia/project/twinkle/cookbook/client/server/transformer/server_config_dsv4_0731_npu_2node_2npu.yaml) +- [4 卡服务端启动脚本](/Users/linjiajia/project/twinkle/cookbook/client/server/transformer/run_dsv4_0731_npu_2node_2npu.sh) +- [4 卡客户端脚本](/Users/linjiajia/project/twinkle/cookbook/client/twinkle/run_dsv4_0731_npu_2node_2npu_client.sh) + +主节点 `172.61.10.111`： + +```bash +DSV4_MODEL_ID=/你的/减层模型路径 \ +NETWORK_IFACE=eth0 \ +RESET_RAY=1 \ +bash cookbook/client/server/transformer/run_dsv4_0731_npu_2node_2npu.sh head +``` + +Worker 节点 `172.61.12.165`： + +```bash +DSV4_MODEL_ID=/你的/减层模型路径 \ +NETWORK_IFACE=eth0 \ +RESET_RAY=1 \ +bash cookbook/client/server/transformer/run_dsv4_0731_npu_2node_2npu.sh worker +``` + +客户端： + +```bash +DSV4_MODEL_ID=/你的/减层模型路径 \ +DATASET_ID=/model/ljl/dataset/self-cognition.jsonl \ +OUTPUT_DIR=/shared/twinkle_output/dsv4-0731-a3-2node-4npu \ +bash cookbook/client/twinkle/run_dsv4_0731_npu_2node_2npu_client.sh +``` + +两个 IP 已写成默认值，通常不需要再传 `HEAD_IP` 或 `NODE_IP`。配置为： + +```text +每节点 NPU：2 +节点数：2 +总 rank：4 +FSDP size：4 +EP size：4 +默认 batch size：4 +``` + +`DSV4_MODEL_ID` 已直接接入 YAML 环境变量解析，因此替换模型时不需要修改 YAML，只需确保主节点、Worker 和客户端传入相同值。语法和配置拓扑校验均已通过。 \ No newline at end of file diff --git a/cookbook/client/server/transformer/run_dsv4_0731_npu_2node_2npu.sh b/cookbook/client/server/transformer/run_dsv4_0731_npu_2node_2npu.sh new file mode 100755 index 000000000..7f3c9404c --- /dev/null +++ b/cookbook/client/server/transformer/run_dsv4_0731_npu_2node_2npu.sh @@ -0,0 +1,161 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Head node (172.61.10.111): +# bash cookbook/client/server/transformer/run_dsv4_0731_npu_2node_2npu.sh head +# +# Worker node (172.61.12.165): +# bash cookbook/client/server/transformer/run_dsv4_0731_npu_2node_2npu.sh worker +# +# Override NETWORK_IFACE when the Ray/HCCL interface is not eth0. Set +# RESET_RAY=1 when intentionally replacing an existing local Ray runtime. + +ROLE="${1:-}" +if [[ "$ROLE" != "head" && "$ROLE" != "worker" ]]; then + echo "Usage: $0 {head|worker}" >&2 + exit 2 +fi + +DEFAULT_HEAD_IP="172.61.10.111" +DEFAULT_WORKER_IP="172.61.12.165" +HEAD_IP="${HEAD_IP:-$DEFAULT_HEAD_IP}" +WORKER_IP="${WORKER_IP:-$DEFAULT_WORKER_IP}" + +if [[ "$ROLE" == "head" ]]; then + NODE_IP="${NODE_IP:-$HEAD_IP}" + NODE_RANK=0 +else + NODE_IP="${NODE_IP:-$WORKER_IP}" + NODE_RANK=1 +fi + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +PROJECT_DIR="$(cd -- "$SCRIPT_DIR/../../../.." && pwd)" +CONFIG_PATH="$SCRIPT_DIR/server_config_dsv4_0731_npu_2node_2npu.yaml" + +export DSV4_MODEL_ID="${DSV4_MODEL_ID:-hf://deepseek-ai/DeepSeek-V4-Flash-0731}" +DATASET_PATH="${DATASET_ID:-/model/ljl/dataset/self-cognition.jsonl}" + +NPU_PER_NODE=2 +NNODES=2 +TOTAL_NPUS=$((NPU_PER_NODE * NNODES)) +RAY_PORT="${RAY_PORT:-6379}" +RAY_CPUS_PER_NODE="${TWINKLE_RAY_CPUS:-$(nproc)}" +CLUSTER_WAIT_SECONDS="${CLUSTER_WAIT_SECONDS:-1800}" +NETWORK_IFACE="${NETWORK_IFACE:-eth0}" + +if [[ -f /usr/local/Ascend/ascend-toolkit/set_env.sh ]]; then + # shellcheck disable=SC1091 + source /usr/local/Ascend/ascend-toolkit/set_env.sh +fi + +cd "$PROJECT_DIR" +export PYTHONPATH="$PROJECT_DIR/src${PYTHONPATH:+:$PYTHONPATH}" +export ASCEND_RT_VISIBLE_DEVICES="${ASCEND_RT_VISIBLE_DEVICES:-0,1}" +export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1 +export TWINKLE_TRUST_REMOTE_CODE=1 +export TWINKLE_FAIL_FAST=1 +export TOKENIZERS_PARALLELISM=true +export GLOO_SOCKET_IFNAME="$NETWORK_IFACE" +export HCCL_SOCKET_IFNAME="$NETWORK_IFACE" +export HCCL_CONNECT_TIMEOUT="${HCCL_CONNECT_TIMEOUT:-7200}" +export HCCL_EXEC_TIMEOUT="${HCCL_EXEC_TIMEOUT:-0}" +export HCCL_IF_BASE_PORT="${HCCL_IF_BASE_PORT:-20000}" +export NODE_RANK + +IFS=',' read -r -a VISIBLE_NPUS <<< "$ASCEND_RT_VISIBLE_DEVICES" +if [[ "${#VISIBLE_NPUS[@]}" -ne "$NPU_PER_NODE" ]]; then + echo "Expected $NPU_PER_NODE visible NPUs, got ASCEND_RT_VISIBLE_DEVICES=$ASCEND_RT_VISIBLE_DEVICES" >&2 + exit 1 +fi + +test -f "$CONFIG_PATH" +test -f "$DATASET_PATH" + +if [[ "$DSV4_MODEL_ID" == hf://* || "$DSV4_MODEL_ID" == ms://* ]]; then + echo "Model will be downloaded through the configured Hub backend: $DSV4_MODEL_ID" +elif [[ -d "$DSV4_MODEL_ID" ]]; then + test -f "$DSV4_MODEL_ID/config.json" + test -f "$DSV4_MODEL_ID/tokenizer.json" +else + echo "Invalid DSV4_MODEL_ID: use an existing local directory or an explicit hf:// or ms:// ID." >&2 + exit 1 +fi + +python3 - "$NPU_PER_NODE" <<'PY' +import sys +import torch +import torch_npu # noqa: F401 + +expected = int(sys.argv[1]) +if not torch.npu.is_available(): + raise SystemExit('torch.npu.is_available() is False') +actual = torch.npu.device_count() +if actual < expected: + raise SystemExit(f'Expected at least {expected} visible NPUs, got {actual}') +print(f'Ascend check passed: visible NPU count={actual}') +PY + +if [[ "${RESET_RAY:-0}" == "1" ]]; then + ray stop --force || true +fi + +if ray status >/dev/null 2>&1; then + echo "Ray is already running on this node; set RESET_RAY=1 to replace it." +elif [[ "$ROLE" == "head" ]]; then + ray start \ + --head \ + --node-ip-address="$HEAD_IP" \ + --port="$RAY_PORT" \ + --num-cpus="$RAY_CPUS_PER_NODE" \ + --resources="{\"NPU\": $NPU_PER_NODE}" \ + --disable-usage-stats \ + --include-dashboard=false +else + ray start \ + --address="$HEAD_IP:$RAY_PORT" \ + --node-ip-address="$NODE_IP" \ + --num-cpus="$RAY_CPUS_PER_NODE" \ + --resources="{\"NPU\": $NPU_PER_NODE}" \ + --disable-usage-stats +fi + +if [[ "$ROLE" == "worker" ]]; then + echo "Worker $NODE_IP joined Ray at $HEAD_IP:$RAY_PORT with $NPU_PER_NODE NPU resources." + exit 0 +fi + +python3 - "$TOTAL_NPUS" "$NNODES" "$NPU_PER_NODE" "$CLUSTER_WAIT_SECONDS" <<'PY' +import sys +import time +import ray + +expected_total = int(sys.argv[1]) +expected_nodes = int(sys.argv[2]) +expected_per_node = int(sys.argv[3]) +timeout = int(sys.argv[4]) +deadline = time.monotonic() + timeout + +ray.init(address='auto', logging_level='ERROR') +try: + while True: + alive_nodes = [node for node in ray.nodes() if node.get('Alive', True)] + npu_nodes = [ + node for node in alive_nodes + if float(node.get('Resources', {}).get('NPU', 0)) >= expected_per_node + ] + total = int(ray.cluster_resources().get('NPU', 0)) + print(f'Waiting for cluster: NPU={total}/{expected_total}, NPU nodes={len(npu_nodes)}/{expected_nodes}') + if total >= expected_total and len(npu_nodes) >= expected_nodes: + break + if time.monotonic() >= deadline: + raise SystemExit( + f'Timed out waiting for {expected_nodes} nodes and {expected_total} total NPU resources') + time.sleep(5) +finally: + ray.shutdown() +PY + +python3 -m twinkle.server check-config --config "$CONFIG_PATH" +echo "Launching Twinkle Server at http://$HEAD_IP:8000" +exec python3 -m twinkle.server launch --config "$CONFIG_PATH" diff --git a/cookbook/client/server/transformer/server_config_dsv4_0731_npu_2node_2npu.yaml b/cookbook/client/server/transformer/server_config_dsv4_0731_npu_2node_2npu.yaml new file mode 100644 index 000000000..f03121f8e --- /dev/null +++ b/cookbook/client/server/transformer/server_config_dsv4_0731_npu_2node_2npu.yaml @@ -0,0 +1,121 @@ +# DeepSeek-V4-compatible Multi-LoRA SFT on two Ascend A3 nodes. +# Each node contributes 2 NPUs, giving a 4-rank Native FSDP2 + EP mesh. +# +# DSV4_MODEL_ID may be either: +# - hf://deepseek-ai/DeepSeek-V4-Flash-0731 +# - ms://namespace/model-name +# - the same absolute local model directory on both nodes + +proxy_location: EveryNode + +http_options: + host: 0.0.0.0 + port: 8000 + +telemetry: + enabled: false + otlp_endpoint: http://localhost:4317 + +persistence: + mode: memory + key_prefix: dsv4-0731-a3-2node-4npu + +applications: + - name: server + route_prefix: /api/v1 + import_path: server + args: + server_config: + per_token_model_limit: 2 + supported_models: + - deepseek-v4-0731-local + deployments: + - name: TinkerCompatServer + max_ongoing_requests: 50 + autoscaling_config: + min_replicas: 1 + max_replicas: 1 + target_ongoing_requests: 32 + ray_actor_options: + num_cpus: 0.1 + runtime_env: + env_vars: + TWINKLE_FAIL_FAST: "1" + + - name: models-deepseek-v4-0731-local + route_prefix: /api/v1/model/deepseek-v4-0731-local + import_path: model + args: + backend: transformers + # Resolved by OmegaConf when the head node reads this YAML. + model_id: ${oc.env:DSV4_MODEL_ID,hf://deepseek-ai/DeepSeek-V4-Flash-0731} + nproc_per_node: 2 + max_loras: 2 + max_r: 8 + max_length: 2048 + mixed_precision: bf16 + strategy: native_fsdp + memory_efficient_init: true + target_modules: all-linear + fsdp_config: + reshard_after_forward: true + expert_parallel: + enabled: true + ep_size: 4 + router_dtype: fp32 + keep_router_logits: false + device_group: + name: model + ranks: 4 + device_type: npu + device_mesh: + device_type: npu + fsdp_size: 4 + dp_size: 1 + ep_size: 4 + queue_config: + rps_limit: 20 + tps_limit: 100000 + max_input_tokens: 4096 + queue_timeout: 7200 + execution_timeout: 7200 + adapter_config: + adapter_timeout: 60 + adapter_max_lifetime: 86400 + deployments: + - name: ModelManagement + autoscaling_config: + min_replicas: 1 + max_replicas: 1 + target_ongoing_requests: 8 + ray_actor_options: + num_cpus: 0.1 + runtime_env: + env_vars: + TWINKLE_TRUST_REMOTE_CODE: "1" + TWINKLE_FAIL_FAST: "1" + + - name: processor + route_prefix: /api/v1/processor + import_path: processor + args: + ncpu_proc_per_node: 2 + device_group: + name: processor + ranks: 2 + device_type: CPU + device_mesh: + device_type: CPU + dp_size: 2 + deployments: + - name: ProcessorManagement + autoscaling_config: + min_replicas: 1 + max_replicas: 1 + target_ongoing_requests: 32 + ray_actor_options: + num_cpus: 0.1 + runtime_env: + env_vars: + TWINKLE_TRUST_REMOTE_CODE: "1" + TWINKLE_FAIL_FAST: "1" diff --git a/cookbook/client/twinkle/run_dsv4_0731_npu_2node_2npu_client.sh b/cookbook/client/twinkle/run_dsv4_0731_npu_2node_2npu_client.sh new file mode 100755 index 000000000..5cce5dfce --- /dev/null +++ b/cookbook/client/twinkle/run_dsv4_0731_npu_2node_2npu_client.sh @@ -0,0 +1,33 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Run exactly one client process. The default server address points at the +# two-node cluster head configured in run_dsv4_0731_npu_2node_2npu.sh. + +HEAD_IP="${HEAD_IP:-172.61.10.111}" +OUTPUT_DIR="${OUTPUT_DIR:-/shared/twinkle_output/dsv4-0731-a3-2node-4npu}" + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +PROJECT_DIR="$(cd -- "$SCRIPT_DIR/../../.." && pwd)" +cd "$PROJECT_DIR" + +export PYTHONPATH="$PROJECT_DIR/src${PYTHONPATH:+:$PYTHONPATH}" +export TWINKLE_SERVER_URL="${TWINKLE_SERVER_URL:-http://$HEAD_IP:8000}" +export TWINKLE_SERVER_TOKEN="${TWINKLE_SERVER_TOKEN:-EMPTY_TOKEN}" +export TWINKLE_MODEL_ID="${TWINKLE_MODEL_ID:-deepseek-v4-0731-local}" +export DSV4_MODEL_ID="${DSV4_MODEL_ID:-hf://deepseek-ai/DeepSeek-V4-Flash-0731}" +export DATASET_ID="${DATASET_ID:-/model/ljl/dataset/self-cognition.jsonl}" +export ADAPTER_NAMES="${ADAPTER_NAMES:-tenant_a,tenant_b}" +export OUTPUT_DIR + +# Four global FSDP ranks: one sample per rank by default. +export BATCH_SIZE="${BATCH_SIZE:-4}" +export GRAD_ACCUM_STEPS="${GRAD_ACCUM_STEPS:-4}" +export MAX_STEPS="${MAX_STEPS:-10}" +export MAX_LENGTH="${MAX_LENGTH:-2048}" +export LORA_R="${LORA_R:-8}" +export LORA_ALPHA="${LORA_ALPHA:-32}" +export LR="${LR:-1e-4}" + +mkdir -p "$OUTPUT_DIR" +exec python3 "$SCRIPT_DIR/dsv4_multi_lora_sft.py" From 9a8e4b272b68fdca9de533fbac4c76516dc5a35a Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Sun, 16 Aug 2026 23:02:43 +0800 Subject: [PATCH 10/28] wip --- src/twinkle/infra/_ray/ray_helper.py | 19 +++++++++++++++++-- tests/infra/test_ray_worker_topology.py | 18 +++++++++++++++++- 2 files changed, 34 insertions(+), 3 deletions(-) diff --git a/src/twinkle/infra/_ray/ray_helper.py b/src/twinkle/infra/_ray/ray_helper.py index c3792812b..3593c1929 100644 --- a/src/twinkle/infra/_ray/ray_helper.py +++ b/src/twinkle/infra/_ray/ray_helper.py @@ -7,6 +7,21 @@ T = TypeVar('T') +# Ray injects these process-local values when it launches each worker. They +# must never be copied from a parent actor into a child runtime_env: a child +# scheduled on another node would otherwise monitor the parent's raylet PID +# and immediately exit with "the local raylet failed". +_RAY_INTERNAL_ENV_VARS = frozenset({ + 'RAY_JOB_ID', + 'RAY_RAYLET_PID', + 'RAY_OVERRIDE_NODE_ID_FOR_TESTING', +}) + + +def _copy_worker_env() -> Dict[str, str]: + """Copy inherited environment without Ray's per-process internal flags.""" + return {key: value for key, value in os.environ.items() if key not in _RAY_INTERNAL_ENV_VARS} + def _get_node_local_topology(placements: List[Dict[str, Any]]) -> List[Tuple[int, List[int]]]: """Return each worker's node-local index and distributed ranks on that node.""" @@ -327,7 +342,7 @@ def create_workers(worker_cls: Type[T], deploy_pg: Dict cluster_name = group worker_name = key + '-' + str(pg_idx) - env_vars = os.environ.copy() + env_vars = _copy_worker_env() env_vars.update({ 'WORLD_SIZE': str(world_size), @@ -395,7 +410,7 @@ def create_workers(worker_cls: Type[T], deploy_pg: Dict cluster_name = group worker_name = key + '-' + str(rank) - env_vars = os.environ.copy() + env_vars = _copy_worker_env() env_vars.update({ 'WORLD_SIZE': str(world_size), 'RANK': str(rank), diff --git a/tests/infra/test_ray_worker_topology.py b/tests/infra/test_ray_worker_topology.py index a4098bd4e..4582a7fb1 100644 --- a/tests/infra/test_ray_worker_topology.py +++ b/tests/infra/test_ray_worker_topology.py @@ -1,4 +1,20 @@ -from twinkle.infra._ray.ray_helper import _get_node_local_topology +from twinkle.infra._ray.ray_helper import _copy_worker_env, _get_node_local_topology + + +def test_copy_worker_env_drops_ray_process_local_flags(monkeypatch): + monkeypatch.setenv('RAY_JOB_ID', 'parent-job') + monkeypatch.setenv('RAY_RAYLET_PID', '12345') + monkeypatch.setenv('RAY_OVERRIDE_NODE_ID_FOR_TESTING', 'parent-node') + monkeypatch.setenv('RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES', '1') + monkeypatch.setenv('TWINKLE_TEST_ENV', 'preserved') + + env = _copy_worker_env() + + assert 'RAY_JOB_ID' not in env + assert 'RAY_RAYLET_PID' not in env + assert 'RAY_OVERRIDE_NODE_ID_FOR_TESTING' not in env + assert env['RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES'] == '1' + assert env['TWINKLE_TEST_ENV'] == 'preserved' def test_get_node_local_topology_for_single_node_workers(): From 0ee30afac43fba28c1ef0924f5422b106f35aab0 Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Tue, 18 Aug 2026 11:27:42 +0800 Subject: [PATCH 11/28] wip --- .../transformer/run_dsv4_0731_npu_multinode.sh | 14 +++++++------- .../server_config_dsv4_0731_npu_multinode.yaml | 3 ++- 2 files changed, 9 insertions(+), 8 deletions(-) diff --git a/cookbook/client/server/transformer/run_dsv4_0731_npu_multinode.sh b/cookbook/client/server/transformer/run_dsv4_0731_npu_multinode.sh index 55a5aaba1..b2547993c 100755 --- a/cookbook/client/server/transformer/run_dsv4_0731_npu_multinode.sh +++ b/cookbook/client/server/transformer/run_dsv4_0731_npu_multinode.sh @@ -25,7 +25,7 @@ fi SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" PROJECT_DIR="$(cd -- "$SCRIPT_DIR/../../../.." && pwd)" CONFIG_PATH="$SCRIPT_DIR/server_config_dsv4_0731_npu_multinode.yaml" -MODEL_ID="hf://deepseek-ai/DeepSeek-V4-Flash-0731" +export DSV4_MODEL_ID="${DSV4_MODEL_ID:-hf://deepseek-ai/DeepSeek-V4-Flash-0731}" DATASET_PATH="${DATASET_ID:-/model/ljl/dataset/self-cognition.jsonl}" NPU_PER_NODE=16 @@ -69,13 +69,13 @@ fi test -f "$CONFIG_PATH" test -f "$DATASET_PATH" -if [[ "$MODEL_ID" == hf://* || "$MODEL_ID" == ms://* ]]; then - echo "Model will be downloaded through the configured Hub backend: $MODEL_ID" -elif [[ -d "$MODEL_ID" ]]; then - test -f "$MODEL_ID/config.json" - test -f "$MODEL_ID/tokenizer.json" +if [[ "$DSV4_MODEL_ID" == hf://* || "$DSV4_MODEL_ID" == ms://* ]]; then + echo "Model will be downloaded through the configured Hub backend: $DSV4_MODEL_ID" +elif [[ -d "$DSV4_MODEL_ID" ]]; then + test -f "$DSV4_MODEL_ID/config.json" + test -f "$DSV4_MODEL_ID/tokenizer.json" else - echo "Invalid MODEL_ID: use an existing local directory or an explicit hf:// or ms:// ID." >&2 + echo "Invalid DSV4_MODEL_ID: use an existing local directory or an explicit hf:// or ms:// ID." >&2 exit 1 fi diff --git a/cookbook/client/server/transformer/server_config_dsv4_0731_npu_multinode.yaml b/cookbook/client/server/transformer/server_config_dsv4_0731_npu_multinode.yaml index 955f98bf6..fee1fc7a8 100644 --- a/cookbook/client/server/transformer/server_config_dsv4_0731_npu_multinode.yaml +++ b/cookbook/client/server/transformer/server_config_dsv4_0731_npu_multinode.yaml @@ -51,7 +51,8 @@ applications: backend: transformers # Twinkle treats an unprefixed ID as ModelScope. Keep the hf:// prefix # here to select Hugging Face explicitly. - model_id: hf://deepseek-ai/DeepSeek-V4-Flash-0731 + # Resolved from the launch environment on the head node. + model_id: ${oc.env:DSV4_MODEL_ID,hf://deepseek-ai/DeepSeek-V4-Flash-0731} nproc_per_node: 16 max_loras: 2 max_r: 8 From 88857df83f4e1fa9f40793a77f1882a73884fdc6 Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Thu, 20 Aug 2026 10:48:45 +0800 Subject: [PATCH 12/28] wip --- .../transformer/server_config_dsv4_0731.yaml | 4 +- ...server_config_dsv4_0731_npu_multinode.yaml | 8 +- .../client/twinkle/dsv4_multi_lora_sft.py | 106 ++++++++++++++---- .../run_dsv4_0731_npu_multinode_client.sh | 14 ++- 4 files changed, 105 insertions(+), 27 deletions(-) diff --git a/cookbook/client/server/transformer/server_config_dsv4_0731.yaml b/cookbook/client/server/transformer/server_config_dsv4_0731.yaml index 2b93293ed..5a8b49542 100644 --- a/cookbook/client/server/transformer/server_config_dsv4_0731.yaml +++ b/cookbook/client/server/transformer/server_config_dsv4_0731.yaml @@ -49,7 +49,7 @@ applications: nproc_per_node: 2 max_loras: 2 max_r: 8 - max_length: 2048 + max_length: 8192 mixed_precision: bf16 strategy: native_fsdp memory_efficient_init: true @@ -73,7 +73,7 @@ applications: queue_config: rps_limit: 20 tps_limit: 100000 - max_input_tokens: 4096 + max_input_tokens: 8192 queue_timeout: 3600 execution_timeout: 3600 adapter_config: diff --git a/cookbook/client/server/transformer/server_config_dsv4_0731_npu_multinode.yaml b/cookbook/client/server/transformer/server_config_dsv4_0731_npu_multinode.yaml index fee1fc7a8..c82324fad 100644 --- a/cookbook/client/server/transformer/server_config_dsv4_0731_npu_multinode.yaml +++ b/cookbook/client/server/transformer/server_config_dsv4_0731_npu_multinode.yaml @@ -28,7 +28,7 @@ applications: import_path: server args: server_config: - per_token_model_limit: 2 + per_token_model_limit: 1 supported_models: - deepseek-v4-0731-local deployments: @@ -54,9 +54,9 @@ applications: # Resolved from the launch environment on the head node. model_id: ${oc.env:DSV4_MODEL_ID,hf://deepseek-ai/DeepSeek-V4-Flash-0731} nproc_per_node: 16 - max_loras: 2 + max_loras: 1 max_r: 8 - max_length: 2048 + max_length: 8192 mixed_precision: bf16 strategy: native_fsdp memory_efficient_init: true @@ -80,7 +80,7 @@ applications: queue_config: rps_limit: 20 tps_limit: 100000 - max_input_tokens: 4096 + max_input_tokens: 8192 queue_timeout: 7200 execution_timeout: 7200 adapter_config: diff --git a/cookbook/client/twinkle/dsv4_multi_lora_sft.py b/cookbook/client/twinkle/dsv4_multi_lora_sft.py index 41f58dfcd..c4f177ff3 100644 --- a/cookbook/client/twinkle/dsv4_multi_lora_sft.py +++ b/cookbook/client/twinkle/dsv4_multi_lora_sft.py @@ -1,10 +1,10 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""Client-side Multi-LoRA SFT for a local DeepSeek-V4-Flash-0731 server. +"""Client-side GSM8K LoRA SFT for a DeepSeek-V4-Flash-0731 server. -Start ``server_config_dsv4_0731.yaml`` first. This client creates two LoRA -tenants on the same two-GPU EP/FSDP model and alternates their SFT micro-steps. -The base model and dataset are both local paths on the server machine. +Start the DeepSeek-V4 multi-node server first. This client trains one LoRA +tenant on the EP/FSDP model. """ +import math import os from peft import LoraConfig @@ -21,28 +21,59 @@ SERVER_TOKEN = os.environ.get('TWINKLE_SERVER_TOKEN', 'EMPTY_TOKEN') SERVED_MODEL_NAME = os.environ.get('TWINKLE_MODEL_ID', 'deepseek-v4-0731-local') MODEL_PATH = os.environ.get('DSV4_MODEL_ID', '/nas/disk1/random-deepseek-v4-4b') -DATASET_PATH = os.environ.get('DATASET_ID', '/model/ljl/dataset/self-cognition.jsonl') +DATASET_ID = os.environ.get('DATASET_ID') +DATASET_SUBSET = os.environ.get('DATASET_SUBSET', 'default') +DATASET_SPLIT = os.environ.get('DATASET_SPLIT', 'train') OUTPUT_DIR = os.environ.get('OUTPUT_DIR', '/tmp/twinkle_dsv4_0731_multi_lora') ADAPTER_NAMES = tuple( - name.strip() for name in os.environ.get('ADAPTER_NAMES', 'tenant_a,tenant_b').split(',') if name.strip()) -BATCH_SIZE = int(os.environ.get('BATCH_SIZE', '2')) + name.strip() for name in os.environ.get('ADAPTER_NAMES', 'tenant_a').split(',') if name.strip()) +BATCH_SIZE = int(os.environ.get('BATCH_SIZE', '32')) GRAD_ACCUM_STEPS = int(os.environ.get('GRAD_ACCUM_STEPS', '4')) MAX_STEPS = int(os.environ.get('MAX_STEPS', '10')) -MAX_LENGTH = int(os.environ.get('MAX_LENGTH', '2048')) +MAX_LENGTH = int(os.environ.get('MAX_LENGTH', '8192')) +TRUNCATION_STRATEGY = os.environ.get('TRUNCATION_STRATEGY', 'delete') LR = float(os.environ.get('LR', '1e-4')) LORA_R = int(os.environ.get('LORA_R', '8')) LORA_ALPHA = int(os.environ.get('LORA_ALPHA', '32')) +def _assert_finite_output(value, path='result') -> None: + """Fail the inference smoke test when a returned numeric value is NaN/Inf.""" + if isinstance(value, dict): + for key, item in value.items(): + _assert_finite_output(item, f'{path}.{key}') + elif isinstance(value, (list, tuple)): + for index, item in enumerate(value): + _assert_finite_output(item, f'{path}[{index}]') + elif isinstance(value, float) and not math.isfinite(value): + raise RuntimeError(f'Inference smoke test produced a non-finite value at {path}: {value}') + + def _build_dataset() -> Dataset: - dataset = Dataset(dataset_meta=DatasetMeta(DATASET_PATH)) - dataset.set_template('DeepseekV4Template', model_id=MODEL_PATH, max_length=MAX_LENGTH) + if not DATASET_ID: + raise ValueError('Set DATASET_ID to a local GSM8K JSON/JSONL file or directory.') + if DATASET_ID.startswith(('hf://', 'ms://')): + raise ValueError(f'DATASET_ID must be local for this recipe, got: {DATASET_ID}') + if not os.path.exists(DATASET_ID): + raise FileNotFoundError(f'Local GSM8K dataset not found: {DATASET_ID}') + dataset = Dataset(dataset_meta=DatasetMeta( + DATASET_ID, + subset_name=DATASET_SUBSET, + split=DATASET_SPLIT, + )) + dataset.set_template( + 'DeepseekV4Template', + model_id=MODEL_PATH, + max_length=MAX_LENGTH, + truncation_strategy=TRUNCATION_STRATEGY, + ) dataset.map( - 'SelfCognitionProcessor', + 'GSM8KProcessor', init_args={ - 'model_name': 'twinkle模型', - 'model_author': 'ModelScope社区', + 'system': 'Solve the math problem step by step and put the final answer in \\boxed{}.', + # SFT needs the reference solution as the assistant target. + 'add_assistant': True, }, ) dataset.encode() @@ -68,13 +99,19 @@ def train() -> None: supported_models = [item.model_name for item in client.get_server_capabilities().supported_models] if SERVED_MODEL_NAME not in supported_models: raise RuntimeError(f'{SERVED_MODEL_NAME!r} is not served; available models: {supported_models}') - if len(ADAPTER_NAMES) != 2: - raise ValueError('This two-slot server example requires exactly two ADAPTER_NAMES.') - if BATCH_SIZE < 2: - raise ValueError('BATCH_SIZE must be at least 2 because the model uses two FSDP data ranks.') + if len(ADAPTER_NAMES) != 1: + raise ValueError('This recipe requires exactly one ADAPTER_NAMES entry.') + if BATCH_SIZE < 32 or BATCH_SIZE % 32 != 0: + raise ValueError('BATCH_SIZE must be at least 32 and divisible by 32 for the 32-rank FSDP model.') dataset = _build_dataset() - dataloaders = {name: DataLoader(dataset=dataset, batch_size=BATCH_SIZE) for name in ADAPTER_NAMES} + # Drop the final undersized batch because 32 FSDP ranks require a full + # global batch here. GSM8K is large enough for the configured run, so each + # adapter traverses at most one shuffled epoch. + dataloaders = { + name: DataLoader(dataset=dataset, batch_size=BATCH_SIZE, drop_last=True, shuffle=True) + for name in ADAPTER_NAMES + } models = {name: MultiLoraTransformersModel(model_id=SERVED_MODEL_NAME) for name in ADAPTER_NAMES} lora_config = _build_lora_config() @@ -89,7 +126,11 @@ def train() -> None: gradient_accumulation_steps=GRAD_ACCUM_STEPS, save_dir=save_dir, ) - model.set_template('DeepseekV4Template', max_length=MAX_LENGTH) + model.set_template( + 'DeepseekV4Template', + max_length=MAX_LENGTH, + truncation_strategy=TRUNCATION_STRATEGY, + ) model.set_processor('InputProcessor', padding_side='right') model.set_loss('CrossEntropyLoss') @@ -127,6 +168,33 @@ def train() -> None: ) logger.info('Saved adapter %s: %s', name, checkpoint.twinkle_path) + # Run a real no-grad forward pass after saving. Compare the trained LoRA + # path with the base-model path so both can be verified on the same input. + # Saving first preserves the expensive training result if this check fails. + eval_loader = DataLoader(dataset=dataset, batch_size=BATCH_SIZE, drop_last=True, shuffle=False) + eval_batch = next(iter(eval_loader)) + for name, model in models.items(): + lora_result = model.forward_only(inputs=eval_batch, disable_lora=False).result + lora_loss = model.calculate_loss().result + base_result = model.forward_only(inputs=eval_batch, disable_lora=True).result + base_loss = model.calculate_loss().result + _assert_finite_output(lora_result, f'{name}.lora') + _assert_finite_output(base_result, f'{name}.base') + _assert_finite_output(lora_loss, f'{name}.lora_loss') + _assert_finite_output(base_loss, f'{name}.base_loss') + # Use WARNING so the inference result remains visible even when the + # deployment intentionally suppresses INFO logs to save local storage. + logger.warning( + 'Inference result: adapter=%s lora_loss=%.6f base_loss=%.6f loss_delta=%.6f ' + 'lora_keys=%s base_keys=%s', + name, + lora_loss, + base_loss, + lora_loss - base_loss, + list(lora_result) if isinstance(lora_result, dict) else type(lora_result).__name__, + list(base_result) if isinstance(base_result, dict) else type(base_result).__name__, + ) + if __name__ == '__main__': train() diff --git a/cookbook/client/twinkle/run_dsv4_0731_npu_multinode_client.sh b/cookbook/client/twinkle/run_dsv4_0731_npu_multinode_client.sh index 1e0ff2006..e56542d48 100755 --- a/cookbook/client/twinkle/run_dsv4_0731_npu_multinode_client.sh +++ b/cookbook/client/twinkle/run_dsv4_0731_npu_multinode_client.sh @@ -20,16 +20,26 @@ export PYTHONPATH="$PROJECT_DIR/src${PYTHONPATH:+:$PYTHONPATH}" export TWINKLE_SERVER_URL="${TWINKLE_SERVER_URL:-http://$HEAD_IP:8000}" export TWINKLE_SERVER_TOKEN="${TWINKLE_SERVER_TOKEN:-EMPTY_TOKEN}" export TWINKLE_MODEL_ID="${TWINKLE_MODEL_ID:-deepseek-v4-0731-local}" +export ADAPTER_NAMES="${ADAPTER_NAMES:-tenant_a}" # The template/encoding side must resolve the same checkpoint as the server. export DSV4_MODEL_ID="${DSV4_MODEL_ID:-hf://deepseek-ai/DeepSeek-V4-Flash-0731}" -export DATASET_ID="${DATASET_ID:-/model/ljl/dataset/self-cognition.jsonl}" +: "${DATASET_ID:?Set DATASET_ID to a local GSM8K JSON/JSONL file or directory}" +if [[ "$DATASET_ID" == hf://* || "$DATASET_ID" == ms://* ]]; then + echo "DATASET_ID must be local for this recipe, got: $DATASET_ID" >&2 + exit 1 +fi +test -e "$DATASET_ID" +export DATASET_ID +export DATASET_SUBSET="${DATASET_SUBSET:-default}" +export DATASET_SPLIT="${DATASET_SPLIT:-train}" export OUTPUT_DIR # forward_backward dispatches one slice to each of the 32 FSDP ranks. export BATCH_SIZE="${BATCH_SIZE:-32}" export GRAD_ACCUM_STEPS="${GRAD_ACCUM_STEPS:-4}" export MAX_STEPS="${MAX_STEPS:-10}" -export MAX_LENGTH="${MAX_LENGTH:-2048}" +export MAX_LENGTH="${MAX_LENGTH:-8192}" +export TRUNCATION_STRATEGY="${TRUNCATION_STRATEGY:-delete}" export LORA_R="${LORA_R:-8}" export LORA_ALPHA="${LORA_ALPHA:-32}" From c0f54e7ec777b10f32cdb71c261f8d3a512f95d2 Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Fri, 21 Aug 2026 10:13:18 +0800 Subject: [PATCH 13/28] fix statedict --- cookbook/client/server/transformer/read_me.md | 2 +- read_me.md | 2 +- .../model/multi_lora_target_parameters.py | 17 ++-- .../transformers/multi_lora_transformers.py | 3 +- .../model/transformers/strategy/accelerate.py | 5 ++ .../transformers/strategy/native_fsdp.py | 84 +++++++++++++++---- .../test_ep_multi_lora_target_parameters.py | 44 ++++++++++ 7 files changed, 130 insertions(+), 27 deletions(-) diff --git a/cookbook/client/server/transformer/read_me.md b/cookbook/client/server/transformer/read_me.md index 57d239056..21e725919 100644 --- a/cookbook/client/server/transformer/read_me.md +++ b/cookbook/client/server/transformer/read_me.md @@ -44,4 +44,4 @@ EP size：4 默认 batch size：4 ``` -`DSV4_MODEL_ID` 已直接接入 YAML 环境变量解析，因此替换模型时不需要修改 YAML，只需确保主节点、Worker 和客户端传入相同值。语法和配置拓扑校验均已通过。 \ No newline at end of file +`DSV4_MODEL_ID` 已直接接入 YAML 环境变量解析，因此替换模型时不需要修改 YAML，只需确保主节点、Worker 和客户端传入相同值。语法和配置拓扑校验均已通过。 diff --git a/read_me.md b/read_me.md index 295b4c3e4..efe55d02e 100644 --- a/read_me.md +++ b/read_me.md @@ -53,4 +53,4 @@ tokenizer.json encoding/encoding_dsv4.py ``` -若不存在 `encoding/encoding_dsv4.py`，才回退到 Twinkle 内置 encoding。相关检查共 `31 passed, 1 skipped`；当前改动尚未 commit/push。 \ No newline at end of file +若不存在 `encoding/encoding_dsv4.py`，才回退到 Twinkle 内置 encoding。相关检查共 `31 passed, 1 skipped`；当前改动尚未 commit/push。 diff --git a/src/twinkle/model/multi_lora_target_parameters.py b/src/twinkle/model/multi_lora_target_parameters.py index 73ee35488..757b21a46 100644 --- a/src/twinkle/model/multi_lora_target_parameters.py +++ b/src/twinkle/model/multi_lora_target_parameters.py @@ -141,10 +141,9 @@ def _write_parameter(parameter: nn.Parameter, value: torch.Tensor) -> None: local_parameter.copy_(value.to(device=local_parameter.device, dtype=local_parameter.dtype)) return if tuple(value.shape) != tuple(parameter.shape): - raise ValueError( - f'Cannot restore target-parameter LoRA with shape {tuple(value.shape)} to distributed ' - f'parameter with global shape {tuple(parameter.shape)} and local shape ' - f'{tuple(local_parameter.shape)}') + raise ValueError(f'Cannot restore target-parameter LoRA with shape {tuple(value.shape)} to distributed ' + f'parameter with global shape {tuple(parameter.shape)} and local shape ' + f'{tuple(local_parameter.shape)}') from torch.distributed.tensor import distribute_tensor distributed = distribute_tensor( value.to(device=parameter.device, dtype=parameter.dtype), @@ -159,16 +158,14 @@ def save_initial_weights(self) -> None: for slot_name, parameter in self.lora_A.items(): tensor = self._read_parameter(parameter) if tensor.is_meta: - raise RuntimeError( - f'Target-parameter LoRA slot {self.record.key}.{slot_name} is still on meta; ' - 'materialize the model before saving its initial weights.') + raise RuntimeError(f'Target-parameter LoRA slot {self.record.key}.{slot_name} is still on meta; ' + 'materialize the model before saving its initial weights.') self._initial_lora_A[slot_name] = tensor.detach().cpu().clone() def reset_slot(self, slot_name: str) -> None: if self.lora_A[slot_name].is_meta or self.lora_B[slot_name].is_meta: - raise RuntimeError( - f'Target-parameter LoRA slot {self.record.key}.{slot_name} is still on meta; ' - 'materialize the model before resetting it.') + raise RuntimeError(f'Target-parameter LoRA slot {self.record.key}.{slot_name} is still on meta; ' + 'materialize the model before resetting it.') with torch.no_grad(): if slot_name not in self._initial_lora_A: nn.init.kaiming_uniform_(self.lora_A[slot_name], a=math.sqrt(5)) diff --git a/src/twinkle/model/transformers/multi_lora_transformers.py b/src/twinkle/model/transformers/multi_lora_transformers.py index dc941bc9a..cdf4da6db 100644 --- a/src/twinkle/model/transformers/multi_lora_transformers.py +++ b/src/twinkle/model/transformers/multi_lora_transformers.py @@ -293,8 +293,9 @@ def get_state_dict(self, **kwargs): return self.multi_adapter.get_state_dict(kwargs.get('adapter_name')) def _get_adapter_state_dict_for_save(self, adapter_name: str) -> dict: + slot_name = self.multi_adapter.find_lora_by_tenant(adapter_name).adapter_name adapter_state = self.multi_adapter.get_state_dict(adapter_name) - return {key: torch_util.to_local_tensor(value).cpu() for key, value in adapter_state.items()} + return self.strategy.gather_adapter_state_dict(self.model, adapter_state, slot_name) @remote_function(collect='first') def save(self, name, output_dir: Optional[str] = None, interval=1, **kwargs): diff --git a/src/twinkle/model/transformers/strategy/accelerate.py b/src/twinkle/model/transformers/strategy/accelerate.py index f74689d11..4e509a980 100644 --- a/src/twinkle/model/transformers/strategy/accelerate.py +++ b/src/twinkle/model/transformers/strategy/accelerate.py @@ -231,6 +231,11 @@ def get_adapter_state_dict(self, model, adapter_name: str) -> dict: del local return state_dict + def gather_adapter_state_dict(self, model, adapter_state: dict, adapter_name: str) -> dict: + """Move an already filtered Multi-LoRA state dict to CPU.""" + from twinkle.utils import torch_util + return {name: torch_util.to_local_tensor(param).cpu() for name, param in adapter_state.items()} + def _is_lora_state_key(name: str) -> bool: return 'lora_A' in name or 'lora_B' in name or 'lora_embedding' in name diff --git a/src/twinkle/model/transformers/strategy/native_fsdp.py b/src/twinkle/model/transformers/strategy/native_fsdp.py index 70b9fd5b3..526e1ef9f 100644 --- a/src/twinkle/model/transformers/strategy/native_fsdp.py +++ b/src/twinkle/model/transformers/strategy/native_fsdp.py @@ -351,16 +351,16 @@ def get_full_state_dict(self, model) -> dict: ep_group = ep_fsdp_mesh['ep'].get_group() ep_world_size = ep_fsdp_mesh['ep'].size() - ep_expert_names = _detect_ep_expert_names(unwrapped) if ep_world_size > 1 else set() + expert_specs = _collect_ep_expert_shard_specs(unwrapped) if ep_world_size > 1 else {} for name, param in unwrapped.named_parameters(): local_full = torch_util.to_local_tensor(param) - if name in ep_expert_names and ep_world_size > 1 and ep_group is not None: + if name in expert_specs and ep_world_size > 1 and ep_group is not None: local_full = local_full.contiguous().to(Platform.get_local_device()) gathered = [torch.empty_like(local_full) for _ in range(ep_world_size)] dist.all_gather(gathered, local_full, group=ep_group) - local_full = torch.cat(gathered, dim=_ep_expert_state_dict_gather_dim(name)) + local_full = _concat_ep_expert_shards(name, gathered, expert_specs[name]) state_dict[name] = local_full.cpu() del gathered, local_full else: @@ -399,7 +399,7 @@ def get_adapter_state_dict(self, model, adapter_name: str) -> dict: ep_group = ep_fsdp_mesh['ep'].get_group() ep_world_size = ep_fsdp_mesh['ep'].size() - ep_expert_names = _detect_ep_expert_names(unwrapped) if ep_world_size > 1 else set() + expert_specs = _collect_ep_expert_shard_specs(unwrapped) if ep_world_size > 1 else {} adapter_suffix = f'.{adapter_name}.' for name, param in unwrapped.named_parameters(): @@ -407,11 +407,11 @@ def get_adapter_state_dict(self, model, adapter_name: str) -> dict: continue local_full = torch_util.to_local_tensor(param) - if name in ep_expert_names and ep_world_size > 1 and ep_group is not None: + if name in expert_specs and ep_world_size > 1 and ep_group is not None: local_full = local_full.contiguous().to(Platform.get_local_device()) gathered = [torch.empty_like(local_full) for _ in range(ep_world_size)] dist.all_gather(gathered, local_full, group=ep_group) - local_full = torch.cat(gathered, dim=_ep_expert_state_dict_gather_dim(name)) + local_full = _concat_ep_expert_shards(name, gathered, expert_specs[name]) state_dict[name] = local_full.cpu() del gathered, local_full else: @@ -420,6 +420,41 @@ def get_adapter_state_dict(self, model, adapter_name: str) -> dict: return state_dict + def gather_adapter_state_dict(self, model, adapter_state: dict, adapter_name: str) -> dict: + """Gather a tenant-filtered Multi-LoRA state dict across the EP group.""" + unwrapped = self.unwrap_model(model) + ep_mesh = self.ep_fsdp_device_mesh + if ep_mesh is None or ep_mesh['ep'].size() <= 1: + return {name: torch_util.to_local_tensor(param).cpu() for name, param in adapter_state.items()} + + ep_group = ep_mesh['ep'].get_group() + slot_suffix = f'.{adapter_name}.' + normalized_specs = { + _strip_peft_base_prefix(name.replace(slot_suffix, '.')): spec + for name, spec in _collect_ep_expert_shard_specs(unwrapped).items() + } + expert_owner_specs = { + name.split('.experts.', 1)[0] + '.experts.': spec + for name, spec in normalized_specs.items() if '.experts.' in name + } + state_dict = {} + for name, param in adapter_state.items(): + local = torch_util.to_local_tensor(param) + canonical_name = _strip_peft_base_prefix(name) + spec = normalized_specs.get(canonical_name) + if spec is None and '.experts.' in canonical_name: + owner = canonical_name.split('.experts.', 1)[0] + '.experts.' + spec = expert_owner_specs.get(owner) + if spec is not None: + local = local.contiguous().to(Platform.get_local_device()) + gathered = [torch.empty_like(local) for _ in range(ep_mesh['ep'].size())] + dist.all_gather(gathered, local, group=ep_group) + local = _concat_ep_expert_shards(name, gathered, spec) + del gathered + state_dict[name] = local.cpu() + del local + return state_dict + def _detect_ep_expert_names(model: nn.Module) -> Set[str]: candidate_names = set() @@ -436,12 +471,21 @@ def _detect_ep_expert_names(model: nn.Module) -> Set[str]: return candidate_names & actual_param_names -def _ep_expert_state_dict_gather_dim(name: str) -> int: +def _ep_expert_state_dict_gather_dim( + name: str, + shape: Optional[tuple] = None, + experts_per_rank: Optional[int] = None, +) -> int: # PEFT ParamWrapper keeps expert LoRA tensors flattened instead of storing # them as [num_experts, ...]: lora_A is [r * num_experts, in] and lora_B is # [out, r * num_experts]. EP therefore owns a contiguous expert block on # dim 0 for A and dim 1 for B. This is still expert sharding, not LoRA rank # parallelism, so the forward pass does not need an EP all-reduce. + # Current target-parameter/3D PEFT tensors keep experts explicitly on dim 0: + # A=[local_experts, r, in], B=[local_experts, out, r]. Detect this before + # applying the legacy flattened-PEFT convention below. + if shape and len(shape) == 3: + return 0 if '_twinkle_lora_' in name: return 0 if 'lora_B' in name: @@ -449,6 +493,18 @@ def _ep_expert_state_dict_gather_dim(name: str) -> int: return 0 +def _concat_ep_expert_shards(name: str, shards: List[torch.Tensor], spec: Dict[str, int]) -> torch.Tensor: + if not shards: + raise ValueError(f'No EP shards collected for {name}.') + local_shape = tuple(shards[0].shape) + gather_dim = _ep_expert_state_dict_gather_dim(name, local_shape, spec['experts_per_rank']) + result = torch.cat(shards, dim=gather_dim) + if local_shape[0] == spec['experts_per_rank'] and result.shape[0] != spec['num_experts']: + raise RuntimeError(f"EP adapter parameter '{name}' reconstructed {result.shape[0]} experts; " + f"expected {spec['num_experts']}.") + return result + + def _build_mp_policy(mixed_precision: str) -> 'MixedPrecisionPolicy': from torch.distributed.fsdp import MixedPrecisionPolicy if mixed_precision == 'bf16': @@ -621,13 +677,11 @@ def _get_local_rank_info() -> tuple[int, int, int, List[int]]: local_rank = _get_node_local_rank() has_twinkle_topology = TWINKLE_NODE_LOCAL_WORLD_SIZE in os.environ if not has_twinkle_topology and 'LOCAL_WORLD_SIZE' not in os.environ and 'LOCAL_SIZE' not in os.environ: - raise RuntimeError( - 'Native FSDP node-local state loading requires Twinkle Ray worker topology, ' - 'LOCAL_WORLD_SIZE, or LOCAL_SIZE.') + raise RuntimeError('Native FSDP node-local state loading requires Twinkle Ray worker topology, ' + 'LOCAL_WORLD_SIZE, or LOCAL_SIZE.') local_world_size = _get_node_local_world_size() raw_node_ranks = os.environ.get(TWINKLE_NODE_RANKS) - if (local_rank < 0 or local_world_size <= 0 - or (not raw_node_ranks and world_size % local_world_size != 0)): + if (local_rank < 0 or local_world_size <= 0 or (not raw_node_ranks and world_size % local_world_size != 0)): raise RuntimeError(f'Invalid local rank topology: rank={rank}, world_size={world_size}, ' f'local_rank={local_rank}, local_world_size={local_world_size}.') if raw_node_ranks: @@ -839,7 +893,7 @@ def _split_for_ep_pre_distribute(model, model_key: str, value: torch.Tensor, ep_ if not matched: return value - shard_dim = _ep_expert_state_dict_gather_dim(model_key) + shard_dim = _ep_expert_state_dict_gather_dim(model_key, tuple(value.shape)) chunk = value.size(shard_dim) // ep_world_size return value.narrow(shard_dim, ep_rank * chunk, chunk).contiguous() @@ -991,7 +1045,9 @@ def _scatter_ep_adapter_tensor(param_name, full_tensor, sharded_param): local_shape = tuple(sharded_param.size()) _, source_dtype = adapter_metadata[param_name] local_tensor = torch.empty(local_shape, device=device_type, dtype=source_dtype) - shard_dim = _ep_expert_state_dict_gather_dim(param_name) + spec = expert_shard_specs.get(param_name) + experts_per_rank = spec['experts_per_rank'] if spec is not None else None + shard_dim = _ep_expert_state_dict_gather_dim(param_name, local_shape, experts_per_rank) local_dim = local_shape[shard_dim] local_tensor = _scatter_ep_tensor_from_source( full_tensor, diff --git a/tests/moe/test_ep_multi_lora_target_parameters.py b/tests/moe/test_ep_multi_lora_target_parameters.py index 3266e2b2d..8012cc87c 100644 --- a/tests/moe/test_ep_multi_lora_target_parameters.py +++ b/tests/moe/test_ep_multi_lora_target_parameters.py @@ -29,6 +29,50 @@ def test_ep_target_parameter_lora_gather_dim_matches_peft_flattening(): "model.layers.0.mlp.experts._twinkle_lora_gate_up_proj.lora_B.lora_0.weight") == 0 +def test_ep_3d_expert_lora_gathers_both_factors_on_expert_dim(): + _ensure_dummy_zmq() + from twinkle.model.transformers.strategy.native_fsdp import ( + _concat_ep_expert_shards, + _ep_expert_state_dict_gather_dim, + ) + + name = "model.layers.0.mlp.experts.base_layer.lora_B.lora_0.weight" + assert _ep_expert_state_dict_gather_dim(name, (8, 4096, 8), 8) == 0 + + shards = [torch.full((8, 2, 1), rank) for rank in range(4)] + full = _concat_ep_expert_shards(name, shards, {"experts_per_rank": 8, "num_experts": 32}) + assert full.shape == (32, 2, 1) + assert torch.equal(full[:, 0, 0], torch.arange(4).repeat_interleave(8)) + + +def test_ep_3d_expert_lora_load_splits_lora_b_on_expert_dim(): + _ensure_dummy_zmq() + from twinkle.model.transformers.strategy.native_fsdp import _split_for_ep_pre_distribute + + class _Box(nn.Module): + pass + + model = _Box() + model.base_model = _Box() + model.base_model.model = _Box() + model.base_model.model.model = _Box() + model.base_model.model.model.layers = nn.ModuleList([_Box()]) + mlp = _Box() + model.base_model.model.model.layers[0].mlp = mlp + mlp._ep_patched = True + mlp.experts = _Box() + mlp.experts.base_layer = _Box() + mlp.experts.base_layer.lora_B = _Box() + mlp.experts.base_layer.lora_B.lora_0 = _Box() + + key = 'base_model.model.model.layers.0.mlp.experts.base_layer.lora_B.lora_0.weight' + full = torch.arange(32).reshape(32, 1, 1).expand(32, 2, 1) + local = _split_for_ep_pre_distribute(model, key, full, ep_world_size=4, ep_rank=2) + + assert local.shape == (8, 2, 1) + assert torch.equal(local[:, 0, 0], torch.arange(16, 24)) + + class _FakeTensorExperts(nn.Module): def __init__(self, *, device="cpu", dtype=torch.float32): From c314b0afcc29583e4690aca82eb0d01d8d4c74f3 Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Thu, 27 Aug 2026 19:33:37 +0800 Subject: [PATCH 14/28] wip --- cookbook/client/twinkle/dsv4_multi_lora_sft.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/cookbook/client/twinkle/dsv4_multi_lora_sft.py b/cookbook/client/twinkle/dsv4_multi_lora_sft.py index c4f177ff3..24438a526 100644 --- a/cookbook/client/twinkle/dsv4_multi_lora_sft.py +++ b/cookbook/client/twinkle/dsv4_multi_lora_sft.py @@ -36,6 +36,10 @@ LR = float(os.environ.get('LR', '1e-4')) LORA_R = int(os.environ.get('LORA_R', '8')) LORA_ALPHA = int(os.environ.get('LORA_ALPHA', '32')) +ROUTED_EXPERT_TARGET_PARAMETERS = [ + 'mlp.experts.gate_up_proj', + 'mlp.experts.down_proj', +] def _assert_finite_output(value, path='result') -> None: @@ -85,12 +89,11 @@ def _build_lora_config() -> LoraConfig: r=LORA_R, lora_alpha=LORA_ALPHA, lora_dropout=0.0, - target_modules='all-linear', - exclude_modules=['o_a_proj'], - target_parameters=[ - 'mlp.experts.gate_up_proj', - 'mlp.experts.down_proj', - ], + # Train only the fused 3D routed-expert parameters. Do not install + # module LoRA on attention, router, shared experts, or the LM head. + target_modules=None, + target_parameters=ROUTED_EXPERT_TARGET_PARAMETERS, + bias='none', ) From b2277191c3a030252ffb3ce80e2aa1333ac5d404 Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Fri, 28 Aug 2026 15:57:16 +0800 Subject: [PATCH 15/28] wip --- .../dsv4_multi_lora_self_cognition_sft.py | 198 ++++++++++++++++++ 1 file changed, 198 insertions(+) create mode 100644 cookbook/client/twinkle/dsv4_multi_lora_self_cognition_sft.py diff --git a/cookbook/client/twinkle/dsv4_multi_lora_self_cognition_sft.py b/cookbook/client/twinkle/dsv4_multi_lora_self_cognition_sft.py new file mode 100644 index 000000000..e52b58aa4 --- /dev/null +++ b/cookbook/client/twinkle/dsv4_multi_lora_self_cognition_sft.py @@ -0,0 +1,198 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Client-side self-cognition SFT for a DeepSeek-V4-Flash-0731 server. + +Start the DeepSeek-V4 multi-node server first. This client trains one LoRA +tenant and updates only the routed-expert gate/up/down parameters. +""" +import math +import os + +from peft import LoraConfig + +from twinkle import get_logger, init_twinkle_client +from twinkle.dataset import DatasetMeta +from twinkle_client.dataloader import DataLoader +from twinkle_client.dataset import Dataset +from twinkle_client.model import MultiLoraTransformersModel + +logger = get_logger() + +SERVER_URL = os.environ.get('TWINKLE_SERVER_URL', 'http://localhost:8000') +SERVER_TOKEN = os.environ.get('TWINKLE_SERVER_TOKEN', 'EMPTY_TOKEN') +SERVED_MODEL_NAME = os.environ.get('TWINKLE_MODEL_ID', 'deepseek-v4-0731-local') +MODEL_PATH = os.environ.get('DSV4_MODEL_ID', '/nas/disk1/random-deepseek-v4-4b') +DATASET_ID = os.environ.get('DATASET_ID') +OUTPUT_DIR = os.environ.get('OUTPUT_DIR', '/tmp/twinkle_dsv4_0731_self_cognition') +MODEL_NAME = os.environ.get('SELF_COGNITION_MODEL_NAME', 'twinkle模型') +MODEL_AUTHOR = os.environ.get('SELF_COGNITION_MODEL_AUTHOR', 'ModelScope社区') + +ADAPTER_NAME = os.environ.get('ADAPTER_NAME', 'tenant_a').strip() +BATCH_SIZE = int(os.environ.get('BATCH_SIZE', '32')) +GRAD_ACCUM_STEPS = int(os.environ.get('GRAD_ACCUM_STEPS', '1')) +NUM_EPOCHS = int(os.environ.get('NUM_EPOCHS', '3')) +MAX_STEPS = int(os.environ.get('MAX_STEPS', '0')) +MAX_LENGTH = int(os.environ.get('MAX_LENGTH', '8192')) +TRUNCATION_STRATEGY = os.environ.get('TRUNCATION_STRATEGY', 'delete') +LR = float(os.environ.get('LR', '1e-4')) +LORA_R = int(os.environ.get('LORA_R', '8')) +LORA_ALPHA = int(os.environ.get('LORA_ALPHA', '32')) +ROUTED_EXPERT_TARGET_PARAMETERS = [ + 'mlp.experts.gate_up_proj', + 'mlp.experts.down_proj', +] + + +def _assert_finite_output(value, path='result') -> None: + if isinstance(value, dict): + for key, item in value.items(): + _assert_finite_output(item, f'{path}.{key}') + elif isinstance(value, (list, tuple)): + for index, item in enumerate(value): + _assert_finite_output(item, f'{path}[{index}]') + elif isinstance(value, float) and not math.isfinite(value): + raise RuntimeError(f'Inference smoke test produced a non-finite value at {path}: {value}') + + +def _build_dataset() -> Dataset: + if not DATASET_ID: + raise ValueError('Set DATASET_ID to a local self-cognition JSON or JSONL file.') + if DATASET_ID.startswith(('hf://', 'ms://')): + raise ValueError(f'DATASET_ID must be local for this recipe, got: {DATASET_ID}') + if not os.path.exists(DATASET_ID): + raise FileNotFoundError(f'Local self-cognition dataset not found: {DATASET_ID}') + + dataset = Dataset(dataset_meta=DatasetMeta(DATASET_ID)) + dataset.set_template( + 'DeepseekV4Template', + model_id=MODEL_PATH, + max_length=MAX_LENGTH, + truncation_strategy=TRUNCATION_STRATEGY, + ) + dataset.map( + 'SelfCognitionProcessor', + init_args={ + 'model_name': MODEL_NAME, + 'model_author': MODEL_AUTHOR, + }, + ) + dataset.encode() + return dataset + + +def _build_lora_config() -> LoraConfig: + return LoraConfig( + r=LORA_R, + lora_alpha=LORA_ALPHA, + lora_dropout=0.0, + target_modules=None, + target_parameters=ROUTED_EXPERT_TARGET_PARAMETERS, + bias='none', + ) + + +def train() -> None: + client = init_twinkle_client(base_url=SERVER_URL, api_key=SERVER_TOKEN) + supported_models = [item.model_name for item in client.get_server_capabilities().supported_models] + if SERVED_MODEL_NAME not in supported_models: + raise RuntimeError(f'{SERVED_MODEL_NAME!r} is not served; available models: {supported_models}') + if not ADAPTER_NAME: + raise ValueError('ADAPTER_NAME must not be empty.') + if BATCH_SIZE < 32 or BATCH_SIZE % 32 != 0: + raise ValueError('BATCH_SIZE must be at least 32 and divisible by 32 for the 32-rank FSDP model.') + if NUM_EPOCHS <= 0: + raise ValueError('NUM_EPOCHS must be greater than zero.') + if MAX_STEPS < 0: + raise ValueError('MAX_STEPS must be zero (unlimited) or a positive integer.') + + dataset = _build_dataset() + dataloader = DataLoader( + dataset=dataset, + batch_size=BATCH_SIZE, + drop_last=True, + shuffle=True, + ) + model = MultiLoraTransformersModel(model_id=SERVED_MODEL_NAME) + save_dir = os.path.join(OUTPUT_DIR, ADAPTER_NAME) + os.makedirs(save_dir, exist_ok=True) + model.add_adapter_to_model( + ADAPTER_NAME, + _build_lora_config(), + gradient_accumulation_steps=GRAD_ACCUM_STEPS, + save_dir=save_dir, + ) + model.set_template( + 'DeepseekV4Template', + max_length=MAX_LENGTH, + truncation_strategy=TRUNCATION_STRATEGY, + ) + model.set_processor('InputProcessor', padding_side='right') + model.set_loss('CrossEntropyLoss') + model.set_optimizer('AdamW', lr=LR, foreach=False) + + completed_steps = 0 + stop_training = False + for epoch in range(NUM_EPOCHS): + logger.info('Starting epoch %s/%s', epoch + 1, NUM_EPOCHS) + for batch in dataloader: + if MAX_STEPS > 0 and completed_steps >= MAX_STEPS: + stop_training = True + break + model.forward_backward(inputs=batch, gradient_accumulation_steps=GRAD_ACCUM_STEPS) + model.clip_grad_and_step(max_grad_norm=1.0, gradient_accumulation_steps=GRAD_ACCUM_STEPS) + completed_steps += 1 + + if completed_steps % GRAD_ACCUM_STEPS == 0: + metric = model.calculate_metric(is_training=True) + logger.info( + 'adapter=%s epoch=%s/%s micro_step=%s metric=%s', + ADAPTER_NAME, + epoch + 1, + NUM_EPOCHS, + completed_steps, + metric.result, + ) + if stop_training: + break + + if completed_steps == 0: + raise RuntimeError( + f'No full batch was produced: dataset is smaller than BATCH_SIZE={BATCH_SIZE}. ' + 'Use a smaller valid global batch or provide more data.' + ) + + checkpoint = model.save( + name=f'dsv4-0731-{ADAPTER_NAME}-self-cognition-final', + save_optimizer=True, + consumed_train_samples=dataloader.get_state()['consumed_train_samples'], + ) + logger.info('Saved adapter %s: %s', ADAPTER_NAME, checkpoint.twinkle_path) + + eval_loader = DataLoader( + dataset=dataset, + batch_size=BATCH_SIZE, + drop_last=True, + shuffle=False, + ) + eval_batch = next(iter(eval_loader)) + lora_result = model.forward_only(inputs=eval_batch, disable_lora=False).result + lora_loss = model.calculate_loss().result + base_result = model.forward_only(inputs=eval_batch, disable_lora=True).result + base_loss = model.calculate_loss().result + _assert_finite_output(lora_result, f'{ADAPTER_NAME}.lora') + _assert_finite_output(base_result, f'{ADAPTER_NAME}.base') + _assert_finite_output(lora_loss, f'{ADAPTER_NAME}.lora_loss') + _assert_finite_output(base_loss, f'{ADAPTER_NAME}.base_loss') + logger.warning( + 'Inference result: adapter=%s lora_loss=%.6f base_loss=%.6f loss_delta=%.6f ' + 'lora_keys=%s base_keys=%s', + ADAPTER_NAME, + lora_loss, + base_loss, + lora_loss - base_loss, + list(lora_result) if isinstance(lora_result, dict) else type(lora_result).__name__, + list(base_result) if isinstance(base_result, dict) else type(base_result).__name__, + ) + + +if __name__ == '__main__': + train() From 2fc2c0f643a4dba0bda243b414849030bb91f7d8 Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Tue, 1 Sep 2026 09:49:48 +0800 Subject: [PATCH 16/28] generate --- src/twinkle/model/base.py | 4 + src/twinkle/model/multi_lora.py | 10 +- .../model/multi_lora_target_parameters.py | 11 +- .../transformers/multi_lora_transformers.py | 27 +++ .../model/transformers/transformers.py | 173 +++++++++++++++++- .../server/model/backends/mock_model.py | 10 + .../model/backends/transformers_model.py | 5 + src/twinkle/server/model/twinkle_handlers.py | 33 ++++ .../model/multi_lora_transformers.py | 27 +++ src/twinkle_client/types/__init__.py | 3 + src/twinkle_client/types/model.py | 23 ++- tests/model/test_transformers_generate.py | 95 ++++++++++ .../integration/test_mock_mode_startup.py | 2 + tests/server/model/test_mock_model.py | 11 ++ 14 files changed, 425 insertions(+), 9 deletions(-) create mode 100644 tests/model/test_transformers_generate.py diff --git a/src/twinkle/model/base.py b/src/twinkle/model/base.py index 8ea00d696..4aaf3fc8b 100644 --- a/src/twinkle/model/base.py +++ b/src/twinkle/model/base.py @@ -30,6 +30,10 @@ def forward(self, *, inputs: Dict[str, Any], **kwargs) -> ModelOutput: def forward_only(self, *, inputs: Dict[str, Any], **kwargs) -> ModelOutput: ... + def generate(self, *, inputs: Dict[str, Any], **kwargs) -> Any: + """Generate completions with the training model when the backend supports it.""" + raise NotImplementedError(f'{type(self).__name__} does not implement generate()') + @abstractmethod def calculate_loss(self, **kwargs) -> float: ... diff --git a/src/twinkle/model/multi_lora.py b/src/twinkle/model/multi_lora.py index 529dbcd65..1104bfa84 100644 --- a/src/twinkle/model/multi_lora.py +++ b/src/twinkle/model/multi_lora.py @@ -173,9 +173,15 @@ def patch_target_parameters(self, module, target_parameters): self.target_parameter_manager.patch(module, target_parameters) @contextmanager - def adapter(self, tenant_adapter_name: str, disable_lora: bool = False): + def adapter(self, + tenant_adapter_name: str, + disable_lora: bool = False, + cache_target_parameters: bool = True): self.activate_adapter(tenant_adapter_name) - with self.target_parameter_manager.adapter(tenant_adapter_name, disable_lora=disable_lora): + with self.target_parameter_manager.adapter( + tenant_adapter_name, + disable_lora=disable_lora, + cache=cache_target_parameters): if disable_lora: # Temporarily disable all adapters while keeping optimizer_group active with self._disable_lora_context(tenant_adapter_name): diff --git a/src/twinkle/model/multi_lora_target_parameters.py b/src/twinkle/model/multi_lora_target_parameters.py index 757b21a46..5b355e248 100644 --- a/src/twinkle/model/multi_lora_target_parameters.py +++ b/src/twinkle/model/multi_lora_target_parameters.py @@ -227,7 +227,7 @@ def get_delta_weight(self, slot_name: str) -> torch.Tensor: return torch.einsum('e o r, e r i -> e i o', weight_B, weight_A) * self.scaling[slot_name] @contextmanager - def activate(self, slot_name: str | None, disable_lora: bool = False): + def activate(self, slot_name: str | None, disable_lora: bool = False, cache: bool = True): if disable_lora or slot_name is None or slot_name not in self.lora_A: yield return @@ -246,7 +246,10 @@ def activate(self, slot_name: str | None, disable_lora: bool = False): ) module.parametrizations[param_name].original.requires_grad_(requires_grad_before) try: - with nn.utils.parametrize.cached(): + if cache: + with nn.utils.parametrize.cached(): + yield + else: yield finally: if not already_parametrized: @@ -412,11 +415,11 @@ def save_initial_weights(self) -> None: wrapper.save_initial_weights() @contextmanager - def adapter(self, tenant_adapter_name: str, disable_lora: bool = False): + def adapter(self, tenant_adapter_name: str, disable_lora: bool = False, cache: bool = True): slot_name = self.tenant_to_slot.get(tenant_adapter_name) with ExitStack() as stack: for wrapper in self.wrappers: - stack.enter_context(wrapper.activate(slot_name, disable_lora=disable_lora)) + stack.enter_context(wrapper.activate(slot_name, disable_lora=disable_lora, cache=cache)) yield def parameters_for_tenant(self, tenant_adapter_name: str) -> list[nn.Parameter]: diff --git a/src/twinkle/model/transformers/multi_lora_transformers.py b/src/twinkle/model/transformers/multi_lora_transformers.py index cdf4da6db..54b828ef9 100644 --- a/src/twinkle/model/transformers/multi_lora_transformers.py +++ b/src/twinkle/model/transformers/multi_lora_transformers.py @@ -191,6 +191,33 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T with self.multi_adapter.adapter(adapter_name, disable_lora=disable_lora): return super().forward_only(inputs=inputs, **kwargs) + @remote_function(dispatch='all', collect='first', sync=True, lazy_collect=False) + def generate(self, + *, + inputs: Union[InputFeature, List[InputFeature], Trajectory, List[Trajectory]], + generation_config: Optional[Dict[str, Any]] = None, + **kwargs): + adapter_name = kwargs.pop('adapter_name', None) + disable_lora = kwargs.pop('disable_lora', False) + self._check_adapter_valid(adapter_name) + # Target-parameter LoRA uses a temporary parametrization while active. + # FSDP must shard the unparametrized model, so finish lazy wrapping + # before entering the adapter context when generate() is the first call. + self._lazy_wrap_model() + # Generation invokes many forwards inside one context. Do not retain a + # parametrized expert weight across FSDP reshard boundaries; recompute + # the routed-expert LoRA delta when each decoder forward accesses it. + with self.multi_adapter.adapter( + adapter_name, + disable_lora=disable_lora, + cache_target_parameters=False): + return super().generate( + inputs=inputs, + adapter_name=adapter_name, + generation_config=generation_config, + **kwargs, + ) + @remote_function(collect='mean') def calculate_loss(self, **kwargs): self._check_adapter_valid(kwargs.get('adapter_name')) diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index 5f7eb251f..2bad9a2b9 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -757,6 +757,171 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T return_outputs['routed_experts'] = recorded_routing return return_outputs + @staticmethod + def _generated_token_ids(sequence: torch.Tensor, prompt_width: int, eos_token_ids, pad_token_id): + """Return completion IDs without batch padding and whether EOS was reached.""" + if eos_token_ids is None: + eos_ids = set() + elif isinstance(eos_token_ids, int): + eos_ids = {eos_token_ids} + else: + eos_ids = {int(token_id) for token_id in eos_token_ids} + + token_ids = [] + stopped = False + for token_id in sequence[prompt_width:].tolist(): + token_id = int(token_id) + if pad_token_id is not None and token_id == pad_token_id and token_id not in eos_ids: + break + token_ids.append(token_id) + if token_id in eos_ids: + stopped = True + break + return token_ids, stopped + + def _prepare_generate_inputs(self, inputs, optimizer_config): + """Encode trajectories and create a left-padded HF generation batch.""" + if isinstance(inputs, dict): + inputs = [inputs] + else: + inputs = list(inputs) + if not inputs: + raise ValueError('inputs empty, check your generate() inputs') + + template = optimizer_config.template + if self._not_encoded(inputs[0]): + assert template is not None, \ + 'Use set_template before passing Trajectory inputs to generate()' + inputs = template.batch_encode(inputs, add_generation_prompt=True) + + if hasattr(self, 'multi_adapter'): + self.multi_adapter.check_length(inputs) + + prompt_token_ids = [] + for item in inputs: + ids = item.get('input_ids') + if ids is None: + raise ValueError("Every generate() input must contain 'input_ids'") + if torch.is_tensor(ids): + ids = ids.detach().cpu().reshape(-1).tolist() + elif isinstance(ids, np.ndarray): + ids = ids.reshape(-1).tolist() + else: + ids = list(ids) + prompt_token_ids.append([int(token_id) for token_id in ids]) + + # The training processor defaults to right padding and may be configured + # for padding-free batches. Decoder-only generation needs a conventional + # left-padded batch, so use a shallow per-call copy without mutating the + # tenant's training processor. + processor = copy(optimizer_config.processor) + assert isinstance(processor, InputProcessor), 'Set InputProcessor correctly before generate()' + processor.padding_side = 'left' + processor.padding_free = False + model_inputs: Dict[str, Any] = processor( + inputs, + sp_strategy=None, + model=self.model, + hf_config=self.hf_config, + enable_sp=False, + ) + for key in ('labels', 'completion_mask', 'length', 'routed_experts'): + model_inputs.pop(key, None) + + # Rebuild position IDs after left padding. Template position IDs describe + # each unpadded sample and therefore cannot be padded with -1 and passed + # unchanged to GenerationMixin. + attention_mask = model_inputs.get('attention_mask') + if torch.is_tensor(attention_mask) and attention_mask.dim() == 2: + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids.masked_fill_(attention_mask == 0, 0) + model_inputs['position_ids'] = position_ids + + return model_inputs, prompt_token_ids, template + + @remote_function(dispatch='all', collect='first', sync=True, lazy_collect=False) + def generate(self, + *, + inputs: Union[InputFeature, List[InputFeature], Trajectory, List[Trajectory]], + generation_config: Optional[Dict[str, Any]] = None, + **kwargs): + """Generate completions directly with the resident Transformers model. + + Unlike ``forward_only()``, this keeps and reuses ``past_key_values`` via + Hugging Face ``GenerationMixin``. All distributed ranks receive the same + request so FSDP/EP collectives remain aligned; only rank 0's JSON-safe + result is returned by the remote-function collector. + """ + adapter_name = kwargs.pop('adapter_name', self._get_default_group()) + optimizer_config = self.optimizer_group[adapter_name] + self._lazy_wrap_model() + if getattr(self, '_enable_sp', False): + raise NotImplementedError('TransformersModel.generate() does not support sequence parallelism; ' + 'start this model with ulysses/sp size 1.') + + self.model.eval() + model_inputs, prompt_token_ids, template = self._prepare_generate_inputs(inputs, optimizer_config) + generate_model = self.strategy.unwrap_model(self.model) + if not hasattr(generate_model, 'generate'): + raise TypeError(f'{type(generate_model).__name__} does not expose Hugging Face generate()') + + gen_kwargs = dict(generation_config or {}) + gen_kwargs.setdefault('max_new_tokens', 128) + gen_kwargs.setdefault('do_sample', False) + gen_kwargs.setdefault('use_cache', True) + gen_kwargs.setdefault('return_dict_in_generate', False) + if gen_kwargs['return_dict_in_generate']: + raise ValueError('return_dict_in_generate=True is not supported by the Twinkle generate API') + + tokenizer = template.tokenizer if template is not None else None + if tokenizer is not None: + if tokenizer.pad_token_id is None and tokenizer.eos_token_id is not None: + gen_kwargs.setdefault('pad_token_id', tokenizer.eos_token_id) + else: + gen_kwargs.setdefault('pad_token_id', tokenizer.pad_token_id) + gen_kwargs.setdefault('eos_token_id', tokenizer.eos_token_id) + + seed = int(gen_kwargs.pop('seed', 0)) + if dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1: + # Ranks may finish at different moments for different prompts. Keep + # every rank participating until the complete distributed group is done. + gen_kwargs.setdefault('synced_gpus', True) + + num_return_sequences = int(gen_kwargs.get('num_return_sequences', 1)) + rng_state = self._get_training_rng_state() + try: + # Identical seeds keep sampled tokens aligned on all FSDP/EP ranks. + # Restore the training RNG afterwards so an evaluation request does + # not perturb dropout or any later stochastic training operation. + Torch.seed_everything(seed) + with torch.no_grad(): + sequences = generate_model.generate(**model_inputs, **gen_kwargs) + finally: + self._set_training_rng_state(rng_state) + + if not torch.is_tensor(sequences): + raise TypeError(f'generate() returned unsupported output type: {type(sequences).__name__}') + expected = len(prompt_token_ids) * num_return_sequences + if sequences.shape[0] != expected: + raise RuntimeError(f'generate() returned {sequences.shape[0]} sequences, expected {expected}') + + prompt_width = int(model_inputs['input_ids'].shape[-1]) + eos_token_ids = gen_kwargs.get('eos_token_id') + pad_token_id = gen_kwargs.get('pad_token_id') + max_new_tokens = int(gen_kwargs['max_new_tokens']) + results = [] + for sequence_index, sequence in enumerate(sequences): + prompt_index = sequence_index // num_return_sequences + token_ids, stopped = self._generated_token_ids(sequence, prompt_width, eos_token_ids, pad_token_id) + text = template.decode(token_ids, skip_special_tokens=True) if template is not None else '' + results.append({ + 'prompt_token_ids': prompt_token_ids[prompt_index], + 'tokens': token_ids, + 'text': text, + 'stop_reason': 'stop' if stopped or len(token_ids) < max_new_tokens else 'length', + }) + return results + @remote_function(collect='mean') def calculate_loss(self, **kwargs): """Calculate loss @@ -1382,8 +1547,8 @@ def _get_training_rng_state(self): state['device_rng_state'] = None return state - def _load_rng_state(self, rng_path): - rng_state = torch.load(rng_path, map_location='cpu', weights_only=False) + @staticmethod + def _set_training_rng_state(rng_state): random.setstate(rng_state['python_rng_state']) np.random.set_state(rng_state['numpy_rng_state']) torch.set_rng_state(rng_state['torch_rng_state']) @@ -1395,6 +1560,10 @@ def _load_rng_state(self, rng_path): if device_module and hasattr(device_module, 'is_available') and device_module.is_available(): device_module.set_rng_state(device_rng_state) + def _load_rng_state(self, rng_path): + rng_state = torch.load(rng_path, map_location='cpu', weights_only=False) + self._set_training_rng_state(rng_state) + def _restore_training_state(self, checkpoint_dir, *, adapter_name=''): trainer_state_path = os.path.join(checkpoint_dir, 'trainer_state.json') with open(trainer_state_path) as f: diff --git a/src/twinkle/server/model/backends/mock_model.py b/src/twinkle/server/model/backends/mock_model.py index 7bf9866cc..323777ee4 100644 --- a/src/twinkle/server/model/backends/mock_model.py +++ b/src/twinkle/server/model/backends/mock_model.py @@ -90,6 +90,16 @@ def forward(self, *, inputs: Any, **kwargs: Any) -> list[dict[str, Any]]: def forward_only(self, *, inputs: Any, **kwargs: Any) -> list[dict[str, Any]]: return self._build_forward_result(inputs, kwargs.get('adapter_name')) + @remote_function(dispatch='all', collect='first', lazy_collect=False) + def generate(self, *, inputs: Any, **kwargs: Any) -> list[dict[str, Any]]: + generated = [1, 2, 3] + return [{ + 'prompt_token_ids': list(item.get('input_ids', [])), + 'tokens': generated, + 'text': 'mock completion', + 'stop_reason': 'length', + } for item in (inputs if isinstance(inputs, list) else [inputs])] + @remote_function() def forward_backward(self, *, inputs: Any, **kwargs: Any) -> list[Any]: loss = float(np.random.default_rng(self._rng_seed).uniform(0.0, 1.0)) diff --git a/src/twinkle/server/model/backends/transformers_model.py b/src/twinkle/server/model/backends/transformers_model.py index 8dc503bb0..82cfcdfdc 100644 --- a/src/twinkle/server/model/backends/transformers_model.py +++ b/src/twinkle/server/model/backends/transformers_model.py @@ -106,6 +106,11 @@ def forward_only(self, *, inputs: InputFeature | list[InputFeature] | Trajectory output = super().forward_only(inputs=inputs, **kwargs) return to_cpu_safe_output(output) + @remote_function(dispatch='all', collect='first', sync=True, lazy_collect=False) + def generate(self, *, inputs: InputFeature | list[InputFeature] | Trajectory | list[Trajectory], **kwargs): + """Generate with the resident Transformers model rather than a sampler deployment.""" + return super().generate(inputs=inputs, **kwargs) + @remote_function(dispatch='slice_dp', collect=collect_tensor_dict) @nccl_safe def forward_backward(self, *, inputs: InputFeature | list[InputFeature] | Trajectory | list[Trajectory], **kwargs): diff --git a/src/twinkle/server/model/twinkle_handlers.py b/src/twinkle/server/model/twinkle_handlers.py index 5bacdd4b0..6d4c63e1d 100644 --- a/src/twinkle/server/model/twinkle_handlers.py +++ b/src/twinkle/server/model/twinkle_handlers.py @@ -152,6 +152,39 @@ async def _task(): task_type='forward_only', )) + @app.post('/twinkle/generate', response_model=types.GenerateResponse) + async def generate( + request: Request, + body: types.GenerateRequest, + self: ModelManagement = Depends(self_fn), + ) -> types.GenerateResponse: + """Generate directly with the Transformers model deployment.""" + token = await self._on_request_start(request) + adapter_name = _get_twinkle_adapter_name(request, body.adapter_name) + + async def _task(): + self.assert_resource_exists(adapter_name) + extra_kwargs = body.model_extra or {} + inputs = _parse_inputs(body.inputs) + ret = self.model.generate( + inputs=inputs, + adapter_name=adapter_name, + generation_config=body.generation_config, + **extra_kwargs, + ) + return {'result': ret} + + inputs_list = body.inputs if isinstance(body.inputs, list) else [body.inputs] + input_tokens = sum(len(inp.get('input_ids', [])) if isinstance(inp, dict) else 0 for inp in inputs_list) + return await run_task( + self.schedule_task_and_wait( + _task, + model_id=adapter_name, + token=token, + input_tokens=input_tokens, + task_type='generate', + )) + @app.post('/twinkle/calculate_loss', response_model=types.CalculateLossResponse) async def calculate_loss( request: Request, diff --git a/src/twinkle_client/model/multi_lora_transformers.py b/src/twinkle_client/model/multi_lora_transformers.py index bf4ef54af..dafbf8511 100644 --- a/src/twinkle_client/model/multi_lora_transformers.py +++ b/src/twinkle_client/model/multi_lora_transformers.py @@ -8,6 +8,7 @@ ClipGradNormResponse, ForwardBackwardResponse, ForwardResponse, + GenerateResponse, GetStateDictResponse, GetTrainConfigsResponse, SaveResponse, @@ -72,6 +73,32 @@ def forward_only(self, inputs: Any, **kwargs) -> ForwardResponse: response.raise_for_status() return ForwardResponse(**response.json()) + def generate( + self, + inputs: Any, + generation_config: Optional[Dict[str, Any]] = None, + *, + timeout: Optional[int] = None, + **kwargs, + ) -> GenerateResponse: + """Generate text directly with the Transformers training model. + + This does not use the sampler deployment. ``generation_config`` is + forwarded to Hugging Face ``generate()``. + """ + response = http_post( + url=f'{self.server_url}/generate', + json_data={ + 'inputs': inputs, + 'adapter_name': self.adapter_name, + 'generation_config': generation_config or {}, + **kwargs, + }, + timeout=timeout, + ) + response.raise_for_status() + return GenerateResponse(**response.json()) + def calculate_loss(self, **kwargs) -> CalculateLossResponse: """Calculate loss from model outputs.""" response = http_post( diff --git a/src/twinkle_client/types/__init__.py b/src/twinkle_client/types/__init__.py index 49673b0e9..20298bffe 100644 --- a/src/twinkle_client/types/__init__.py +++ b/src/twinkle_client/types/__init__.py @@ -20,6 +20,9 @@ ForwardOnlyRequest, ForwardRequest, ForwardResponse, + GenerateRequest, + GenerateResponse, + GeneratedSequence, GetStateDictRequest, GetStateDictResponse, GetTrainConfigsResponse, diff --git a/src/twinkle_client/types/model.py b/src/twinkle_client/types/model.py index 83f489cde..dc19458f6 100644 --- a/src/twinkle_client/types/model.py +++ b/src/twinkle_client/types/model.py @@ -4,7 +4,7 @@ These models are used by both the server-side handler and the twinkle client. """ -from pydantic import BaseModel, field_validator +from pydantic import BaseModel, Field, field_validator from typing import Any, Dict, List, Optional, Union @@ -30,6 +30,15 @@ class Config: extra = 'allow' +class GenerateRequest(BaseModel): + inputs: Any + adapter_name: Optional[str] = None + generation_config: Dict[str, Any] = Field(default_factory=dict) + + class Config: + extra = 'allow' + + class AdapterRequest(BaseModel): adapter_name: str @@ -202,6 +211,18 @@ class ForwardBackwardResponse(BaseModel): result: Any +class GeneratedSequence(BaseModel): + """One completion returned by the Transformers training model.""" + prompt_token_ids: List[int] + tokens: List[int] + text: str + stop_reason: str + + +class GenerateResponse(BaseModel): + result: List[GeneratedSequence] + + class CalculateLossResponse(BaseModel): """Response for /calculate_loss endpoint (returns float).""" result: float diff --git a/tests/model/test_transformers_generate.py b/tests/model/test_transformers_generate.py new file mode 100644 index 000000000..ccbc6697b --- /dev/null +++ b/tests/model/test_transformers_generate.py @@ -0,0 +1,95 @@ +from types import SimpleNamespace +from unittest.mock import patch + +import numpy as np +import torch + +from twinkle.model.transformers.transformers import TransformersModel +from twinkle.processor import InputProcessor + + +class _Tokenizer: + pad_token_id = 0 + eos_token_id = 2 + + +class _Template: + tokenizer = _Tokenizer() + + def batch_encode(self, rows, add_generation_prompt=False): + assert add_generation_prompt + return [ + { + 'input_ids': np.array([10, 11]), + 'attention_mask': np.ones(2), + 'position_ids': np.arange(2), + 'labels': np.array([-100, -100]), + } for _ in rows + ] + + def decode(self, token_ids, **kwargs): + return ' '.join(map(str, token_ids)) + + +class _GenerateModel: + + def eval(self): + return self + + def generate(self, input_ids, **kwargs): + completion = torch.tensor([[20, 2]], device=input_ids.device).expand(input_ids.shape[0], -1) + return torch.cat([input_ids, completion], dim=-1) + + +class _Strategy: + + @staticmethod + def unwrap_model(model): + return model + + +def _model_wrapper(): + wrapper = object.__new__(TransformersModel) + wrapper.optimizer_group = { + 'tenant': SimpleNamespace(template=_Template(), processor=InputProcessor()), + } + wrapper.model = _GenerateModel() + wrapper.strategy = _Strategy() + wrapper.hf_config = SimpleNamespace() + wrapper._model_wrapped = True + wrapper._enable_sp = False + wrapper._lazy_wrap_model = lambda: None + wrapper._get_default_group = lambda: 'tenant' + return wrapper + + +def test_generate_encodes_trajectory_and_returns_completion_only(): + wrapper = _model_wrapper() + torch.manual_seed(1234) + rng_state = torch.get_rng_state().clone() + with patch('twinkle.processor.base.Platform.get_local_device', return_value=torch.device('cpu')): + result = TransformersModel.generate.__wrapped__( + wrapper, + inputs=[{'messages': [{'role': 'user', 'content': 'hello'}]}], + adapter_name='tenant', + generation_config={'max_new_tokens': 2}, + ) + + assert result == [{ + 'prompt_token_ids': [10, 11], + 'tokens': [20, 2], + 'text': '20 2', + 'stop_reason': 'stop', + }] + assert torch.equal(torch.get_rng_state(), rng_state) + + +def test_generated_token_ids_trims_padding_after_eos(): + tokens, stopped = TransformersModel._generated_token_ids( + torch.tensor([10, 11, 20, 2, 0, 0]), + prompt_width=2, + eos_token_ids=[2], + pad_token_id=0, + ) + assert tokens == [20, 2] + assert stopped is True diff --git a/tests/server/integration/test_mock_mode_startup.py b/tests/server/integration/test_mock_mode_startup.py index be8610ced..f26a2e89b 100644 --- a/tests/server/integration/test_mock_mode_startup.py +++ b/tests/server/integration/test_mock_mode_startup.py @@ -240,6 +240,8 @@ def _exercise_twinkle_clients(base: str) -> None: assert fwd.result is not None fwd_only = model.forward_only(inputs) assert fwd_only.result is not None + generated = model.generate(inputs, generation_config={'max_new_tokens': 3}) + assert generated.result[0].text == 'mock completion' fwd_bwd = model.forward_backward(inputs) assert fwd_bwd.result is not None diff --git a/tests/server/model/test_mock_model.py b/tests/server/model/test_mock_model.py index 5c24450f4..2d26419a2 100644 --- a/tests/server/model/test_mock_model.py +++ b/tests/server/model/test_mock_model.py @@ -29,6 +29,7 @@ 'tinker_calculate_metric', 'tinker_load', 'forward_only', + 'generate', 'forward_backward', 'forward', 'calculate_loss', @@ -88,6 +89,16 @@ def test_forward_only_deterministic_and_shaped(seq_lens: list, seed: int) -> Non assert len(record['elementwise_loss']) == n +def test_generate_returns_json_safe_completion() -> None: + model = TwinkleCompatMockModel('mid') + assert model.generate(inputs=[{'input_ids': [10, 11]}]) == [{ + 'prompt_token_ids': [10, 11], + 'tokens': [1, 2, 3], + 'text': 'mock completion', + 'stop_reason': 'length', + }] + + @settings(max_examples=100) @given(seq_lens=st.lists(st.integers(min_value=1, max_value=8), min_size=1, max_size=4)) def test_tinker_forward_backward_loss_is_finite(seq_lens: list) -> None: From beee992ab4b8d426070f9fb505f2291b33064373 Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Tue, 1 Sep 2026 10:10:13 +0800 Subject: [PATCH 17/28] wip --- .../transformers/strategy/native_fsdp.py | 10 ++++- .../test_native_fsdp_single_device.py | 45 +++++++++++++++++++ 2 files changed, 53 insertions(+), 2 deletions(-) create mode 100644 tests/transformers/test_native_fsdp_single_device.py diff --git a/src/twinkle/model/transformers/strategy/native_fsdp.py b/src/twinkle/model/transformers/strategy/native_fsdp.py index 526e1ef9f..8355cfd42 100644 --- a/src/twinkle/model/transformers/strategy/native_fsdp.py +++ b/src/twinkle/model/transformers/strategy/native_fsdp.py @@ -131,9 +131,15 @@ 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): - if self.device_mesh is None: + fsdp_mesh = _build_fsdp_mesh(self.device_mesh) if self.device_mesh is not None else None + if fsdp_mesh is None: + # FSDP normally materializes/moves parameters onto the mesh device + # while wrapping. A singleton (or absent) mesh skips FSDP, so do + # the equivalent device placement explicitly. Without this, a + # model loaded by ``from_pretrained`` remains on CPU while the + # input processor creates tensors on the actor's CUDA/NPU device. + model = model.to(torch.device(Platform.get_local_device())) return model, optimizer - fsdp_mesh = _build_fsdp_mesh(self.device_mesh) if fsdp_mesh is not None: ep_enabled = (self.enable_ep and self.ep_fsdp_device_mesh is not None) diff --git a/tests/transformers/test_native_fsdp_single_device.py b/tests/transformers/test_native_fsdp_single_device.py new file mode 100644 index 000000000..26972b69a --- /dev/null +++ b/tests/transformers/test_native_fsdp_single_device.py @@ -0,0 +1,45 @@ +import numpy as np +import torch +from torch import nn + +from twinkle import DeviceMesh +from twinkle.model.transformers.strategy.native_fsdp import NativeFSDPStrategy +from twinkle.utils import Platform + + +class _DeviceTrackingModel(nn.Module): + + def __init__(self): + super().__init__() + self.linear = nn.Linear(2, 2) + self.to_devices = [] + + def to(self, *args, **kwargs): + device = args[0] if args else kwargs.get('device') + self.to_devices.append(torch.device(device)) + return super().to(*args, **kwargs) + + +def test_native_fsdp_singleton_mesh_places_model_on_local_device(monkeypatch): + monkeypatch.setattr(Platform, 'get_local_device', lambda: 'cpu') + device_mesh = DeviceMesh( + device_type='cuda', + mesh=np.array([0]), + mesh_dim_names=('fsdp', ), + ) + strategy = NativeFSDPStrategy( + device_mesh=device_mesh, + memory_efficient_init=True, + enable_ep=False, + ) + model = _DeviceTrackingModel() + optimizer = torch.optim.SGD(model.parameters(), lr=0.1) + original_optimizer_params = list(optimizer.param_groups[0]['params']) + + wrapped_model, wrapped_optimizer = strategy.wrap_model(model, optimizer) + + assert wrapped_model is model + assert wrapped_optimizer is optimizer + assert model.to_devices == [torch.device('cpu')] + assert all(param.device.type == 'cpu' for param in model.parameters()) + assert list(optimizer.param_groups[0]['params']) == original_optimizer_params From 8aa9102072d1f41e21eebe29bbb5b56f1bd944b0 Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Tue, 1 Sep 2026 10:28:47 +0800 Subject: [PATCH 18/28] generate --- .../transformers/strategy/native_fsdp.py | 23 +++++++++++++ .../model/transformers/transformers.py | 5 ++- tests/model/test_transformers_generate.py | 16 +++++++++ .../test_native_fsdp_single_device.py | 34 +++++++++++++++++++ 4 files changed, 77 insertions(+), 1 deletion(-) diff --git a/src/twinkle/model/transformers/strategy/native_fsdp.py b/src/twinkle/model/transformers/strategy/native_fsdp.py index 8355cfd42..2671284e7 100644 --- a/src/twinkle/model/transformers/strategy/native_fsdp.py +++ b/src/twinkle/model/transformers/strategy/native_fsdp.py @@ -2,6 +2,7 @@ import os import torch import torch.distributed as dist +from contextlib import contextmanager from torch import nn from torch.distributed.device_mesh import DeviceMesh as TorchDeviceMesh from torch.distributed.fsdp import fully_shard @@ -264,6 +265,28 @@ def wrap_model(self, model, optimizer=None): return model, optimizer + @contextmanager + def generation_context(self, model): + """Materialize root FSDP parameters while PEFT delegates generation. + + ``PeftModel.generate()`` calls ``get_base_model().generate()`` directly, + bypassing the outer PEFT module's forward hooks. When that outer module + is an FSDP2 root, its directly managed parameters (for example token + embeddings and ``lm_head``) would otherwise remain DTensors while the + generation inputs are regular tensors. + """ + unshard = getattr(model, 'unshard', None) + reshard = getattr(model, 'reshard', None) + if not callable(unshard) or not callable(reshard): + yield + return + + unshard() + try: + yield + finally: + reshard() + def _prepare_optimizer_state_dict_options(self, *, for_load: bool): from torch.distributed.checkpoint.state_dict import StateDictOptions diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index 2bad9a2b9..34b18ba8c 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -894,7 +894,10 @@ def generate(self, # Restore the training RNG afterwards so an evaluation request does # not perturb dropout or any later stochastic training operation. Torch.seed_everything(seed) - with torch.no_grad(): + generation_context = getattr(self.strategy, 'generation_context', None) + fsdp_root_context = ( + generation_context(generate_model) if generation_context is not None else contextlib.nullcontext()) + with torch.no_grad(), fsdp_root_context: sequences = generate_model.generate(**model_inputs, **gen_kwargs) finally: self._set_training_rng_state(rng_state) diff --git a/tests/model/test_transformers_generate.py b/tests/model/test_transformers_generate.py index ccbc6697b..a4ada6687 100644 --- a/tests/model/test_transformers_generate.py +++ b/tests/model/test_transformers_generate.py @@ -1,3 +1,4 @@ +from contextlib import contextmanager from types import SimpleNamespace from unittest.mock import patch @@ -43,10 +44,21 @@ def generate(self, input_ids, **kwargs): class _Strategy: + def __init__(self): + self.generation_events = [] + @staticmethod def unwrap_model(model): return model + @contextmanager + def generation_context(self, model): + self.generation_events.append(('enter', model)) + try: + yield + finally: + self.generation_events.append(('exit', model)) + def _model_wrapper(): wrapper = object.__new__(TransformersModel) @@ -81,6 +93,10 @@ def test_generate_encodes_trajectory_and_returns_completion_only(): 'text': '20 2', 'stop_reason': 'stop', }] + assert wrapper.strategy.generation_events == [ + ('enter', wrapper.model), + ('exit', wrapper.model), + ] assert torch.equal(torch.get_rng_state(), rng_state) diff --git a/tests/transformers/test_native_fsdp_single_device.py b/tests/transformers/test_native_fsdp_single_device.py index 26972b69a..67ea1c07b 100644 --- a/tests/transformers/test_native_fsdp_single_device.py +++ b/tests/transformers/test_native_fsdp_single_device.py @@ -1,4 +1,5 @@ import numpy as np +import pytest import torch from torch import nn @@ -20,6 +21,18 @@ def to(self, *args, **kwargs): return super().to(*args, **kwargs) +class _FSDPRoot: + + def __init__(self): + self.events = [] + + def unshard(self): + self.events.append('unshard') + + def reshard(self): + self.events.append('reshard') + + def test_native_fsdp_singleton_mesh_places_model_on_local_device(monkeypatch): monkeypatch.setattr(Platform, 'get_local_device', lambda: 'cpu') device_mesh = DeviceMesh( @@ -43,3 +56,24 @@ def test_native_fsdp_singleton_mesh_places_model_on_local_device(monkeypatch): assert model.to_devices == [torch.device('cpu')] assert all(param.device.type == 'cpu' for param in model.parameters()) assert list(optimizer.param_groups[0]['params']) == original_optimizer_params + + +def test_native_fsdp_generation_context_unshards_and_restores_root(): + strategy = NativeFSDPStrategy(device_mesh=None, enable_ep=False) + model = _FSDPRoot() + + with strategy.generation_context(model): + model.events.append('generate') + + assert model.events == ['unshard', 'generate', 'reshard'] + + +def test_native_fsdp_generation_context_reshards_after_failure(): + strategy = NativeFSDPStrategy(device_mesh=None, enable_ep=False) + model = _FSDPRoot() + + with pytest.raises(RuntimeError, match='generation failed'): + with strategy.generation_context(model): + raise RuntimeError('generation failed') + + assert model.events == ['unshard', 'reshard'] From 27226b5f213890a4803b5fa4c789defc08accf16 Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Tue, 1 Sep 2026 18:13:27 +0800 Subject: [PATCH 19/28] wip --- .../diagnostics/compare_dsv4_4layer_logits.py | 70 +++ .../diagnostics/probe_dsv4_4layer_logits.py | 146 +++++ .../diagnostics/test_dsv4_npu_gmm_layout.py | 131 +++++ cookbook/client/server/transformer/read_me.md | 555 +++++++++++++++++- .../run_dsv4_0731_npu_2node_2npu.sh | 13 +- .../run_dsv4_4layer_ep_diagnostic.sh | 37 ++ .../server_config_dsv4_4layer_diag_ep.yaml | 117 ++++ .../server_config_dsv4_4layer_diag_no_ep.yaml | 117 ++++ src/twinkle/kernel/ops/ep/__init__.py | 14 + src/twinkle/kernel/ops/moe/npu.py | 32 +- .../model/transformers/moe/expert_parallel.py | 49 ++ .../transformers/strategy/native_fsdp.py | 32 + tests/kernel/ops/test_moe.py | 67 ++- 13 files changed, 1339 insertions(+), 41 deletions(-) create mode 100755 cookbook/client/server/transformer/diagnostics/compare_dsv4_4layer_logits.py create mode 100755 cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits.py create mode 100755 cookbook/client/server/transformer/diagnostics/test_dsv4_npu_gmm_layout.py create mode 100755 cookbook/client/server/transformer/run_dsv4_4layer_ep_diagnostic.sh create mode 100644 cookbook/client/server/transformer/server_config_dsv4_4layer_diag_ep.yaml create mode 100644 cookbook/client/server/transformer/server_config_dsv4_4layer_diag_no_ep.yaml diff --git a/cookbook/client/server/transformer/diagnostics/compare_dsv4_4layer_logits.py b/cookbook/client/server/transformer/diagnostics/compare_dsv4_4layer_logits.py new file mode 100755 index 000000000..58facf908 --- /dev/null +++ b/cookbook/client/server/transformer/diagnostics/compare_dsv4_4layer_logits.py @@ -0,0 +1,70 @@ +#!/usr/bin/env python3 +"""Numerically compare two last-token-logit files produced by the DSV4 probe.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import torch + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('reference', type=Path) + parser.add_argument('candidate', type=Path) + parser.add_argument('--rtol', type=float, default=1e-2) + parser.add_argument('--atol', type=float, default=2e-2) + parser.add_argument('--output', type=Path) + return parser.parse_args() + + +def load_logits(path: Path) -> tuple[str, list[int], torch.Tensor]: + payload = torch.load(path, map_location='cpu', weights_only=True) + return str(payload.get('mode', path.stem)), list(payload['input_ids']), payload['last_logits'].float() + + +def main() -> None: + args = parse_args() + reference_mode, reference_ids, reference = load_logits(args.reference) + candidate_mode, candidate_ids, candidate = load_logits(args.candidate) + if reference_ids != candidate_ids: + raise RuntimeError(f'Input IDs differ: {reference_ids} != {candidate_ids}') + if reference.shape != candidate.shape: + raise RuntimeError(f'Logit shapes differ: {tuple(reference.shape)} != {tuple(candidate.shape)}') + + difference = (candidate - reference).abs() + close = torch.isclose(candidate, reference, rtol=args.rtol, atol=args.atol) + reference_top = torch.topk(reference, k=min(20, reference.numel())).indices.tolist() + candidate_top = torch.topk(candidate, k=min(20, candidate.numel())).indices.tolist() + report = { + 'reference': str(args.reference), + 'reference_mode': reference_mode, + 'candidate': str(args.candidate), + 'candidate_mode': candidate_mode, + 'input_ids': reference_ids, + 'shape': list(reference.shape), + 'rtol': args.rtol, + 'atol': args.atol, + 'allclose': bool(close.all().item()), + 'close_fraction': close.float().mean().item(), + 'max_abs_diff': difference.max().item(), + 'mean_abs_diff': difference.mean().item(), + 'reference_top20': reference_top, + 'candidate_top20': candidate_top, + 'top20_overlap': len(set(reference_top) & set(candidate_top)), + } + + output = args.output or args.candidate.with_name( + f'compare_{reference_mode}_vs_{candidate_mode}.json') + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding='utf-8') + print(json.dumps(report, ensure_ascii=False, indent=2)) + print(f'Report saved to: {output.resolve()}') + if not report['allclose']: + raise SystemExit(1) + + +if __name__ == '__main__': + main() diff --git a/cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits.py b/cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits.py new file mode 100755 index 000000000..dd40d122d --- /dev/null +++ b/cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits.py @@ -0,0 +1,146 @@ +#!/usr/bin/env python3 +"""Save last-token logits from a running four-layer Twinkle server.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from pathlib import Path +from typing import Any + +import torch +from peft import LoraConfig + +from twinkle_client import init_twinkle_client +from twinkle_client.model import MultiLoraTransformersModel + + +DEFAULT_INPUT_IDS = [0, 128803, 2788, 6573, 70979, 36005, 320, 128804, 128821] +TARGET_PARAMETERS = ['mlp.experts.gate_up_proj', 'mlp.experts.down_proj'] + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--mode', required=True, choices=('no_ep', 'ep_loop', 'ep_gmm')) + parser.add_argument('--server-url', default='http://127.0.0.1:8000') + parser.add_argument('--server-token', default='EMPTY_TOKEN') + parser.add_argument('--served-model', default='deepseek-v4-0731-local') + parser.add_argument('--output-dir', type=Path, default=Path('output/dsv4_ep_diag')) + parser.add_argument('--input-ids', default=','.join(str(item) for item in DEFAULT_INPUT_IDS)) + return parser.parse_args() + + +def _extract_logits(result: Any) -> torch.Tensor: + if hasattr(result, 'model_dump'): + result = result.model_dump() + if isinstance(result, list) and len(result) == 1 and isinstance(result[0], dict): + result = result[0] + if not isinstance(result, dict) or result.get('logits') is None: + raise RuntimeError(f'forward_only did not return logits; result type={type(result).__name__}') + + logits = torch.as_tensor(result['logits'], dtype=torch.float32) + original_shape = tuple(logits.shape) + while logits.ndim > 3: + logits = logits[0] + if logits.ndim == 3: + logits = logits[0, -1] + elif logits.ndim == 2: + logits = logits[-1] + elif logits.ndim != 1: + raise RuntimeError(f'Unsupported logits shape: {original_shape}') + if logits.numel() < 1000: + raise RuntimeError(f'Last-token logits look too small: shape={tuple(logits.shape)}') + return logits.contiguous() + + +def _tensor_sha256(tensor: torch.Tensor) -> str: + values = tensor.detach().cpu().contiguous().numpy().tobytes() + return hashlib.sha256(values).hexdigest() + + +def main() -> None: + args = parse_args() + input_ids = [int(item.strip()) for item in args.input_ids.split(',') if item.strip()] + if not input_ids: + raise SystemExit('--input-ids must contain at least one token ID') + + client = init_twinkle_client( + base_url=args.server_url, + api_key=args.server_token, + session_heartbeat_interval=10, + ) + try: + capacity = client.get_capacity_info() + if capacity.free_loras < 1: + raise RuntimeError('No free LoRA slot. Restart the diagnostic server before running the probe.') + + model = MultiLoraTransformersModel(model_id=args.served_model) + model.add_adapter_to_model( + f'dsv4_diag_{args.mode}', + LoraConfig( + r=8, + lora_alpha=32, + lora_dropout=0.0, + target_modules=None, + target_parameters=TARGET_PARAMETERS, + bias='none', + ), + gradient_accumulation_steps=1, + ) + model.set_processor('InputProcessor', padding_side='left', padding_free=False) + + raw_input = { + 'input_ids': input_ids, + 'attention_mask': [1] * len(input_ids), + 'position_ids': list(range(len(input_ids))), + } + response = model.forward_only( + inputs=[raw_input], + disable_lora=True, + return_logits=True, + ) + last_logits = _extract_logits(response.result) + finally: + client.close() + + finite = torch.isfinite(last_logits) + top_values, top_indices = torch.topk(last_logits, k=min(20, last_logits.numel())) + report = { + 'mode': args.mode, + 'server_url': args.server_url, + 'served_model': args.served_model, + 'input_ids': input_ids, + 'last_logits_shape': list(last_logits.shape), + 'dtype_saved': str(last_logits.dtype), + 'sha256': _tensor_sha256(last_logits), + 'finite': bool(finite.all().item()), + 'nan_count': int(torch.isnan(last_logits).sum().item()), + 'inf_count': int(torch.isinf(last_logits).sum().item()), + 'sum': last_logits.sum().item(), + 'abs_sum': last_logits.abs().sum().item(), + 'min': last_logits.min().item(), + 'max': last_logits.max().item(), + 'top_token_ids': top_indices.tolist(), + 'top_logits': top_values.tolist(), + } + + args.output_dir.mkdir(parents=True, exist_ok=True) + tensor_path = args.output_dir / f'{args.mode}_last_logits.pt' + json_path = args.output_dir / f'{args.mode}_last_logits.json' + torch.save( + { + 'mode': args.mode, + 'input_ids': input_ids, + 'last_logits': last_logits, + }, + tensor_path, + ) + json_path.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding='utf-8') + print(json.dumps(report, ensure_ascii=False, indent=2)) + print(f'Logits saved to: {tensor_path.resolve()}') + print(f'Report saved to: {json_path.resolve()}') + + +if __name__ == '__main__': + main() diff --git a/cookbook/client/server/transformer/diagnostics/test_dsv4_npu_gmm_layout.py b/cookbook/client/server/transformer/diagnostics/test_dsv4_npu_gmm_layout.py new file mode 100755 index 000000000..bd65df18d --- /dev/null +++ b/cookbook/client/server/transformer/diagnostics/test_dsv4_npu_gmm_layout.py @@ -0,0 +1,131 @@ +#!/usr/bin/env python3 +"""Compare the DeepSeek-V4 square expert layout on NPU GMM against F.linear.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import torch +import torch.nn.functional as F +from torch import nn + +from twinkle.kernel.ops.moe.npu import GmmFunction, _normalize_packed_expert_weights + + +class PackedExperts(nn.Module): + + def __init__(self, gate_up_proj: torch.Tensor, down_proj: torch.Tensor): + super().__init__() + self.gate_up_proj = nn.Parameter(gate_up_proj, requires_grad=False) + self.down_proj = nn.Parameter(down_proj, requires_grad=False) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--device', default='npu:0') + parser.add_argument('--dtype', choices=('float16', 'bfloat16'), default='bfloat16') + parser.add_argument('--atol', type=float, default=2e-2) + parser.add_argument('--rtol', type=float, default=1e-2) + parser.add_argument('--output', type=Path, default=Path('output/dsv4_ep_diag/npu_gmm_layout.json')) + return parser.parse_args() + + +def reference_forward( + inputs: torch.Tensor, + counts: list[int], + gate_up_proj: torch.Tensor, + down_proj: torch.Tensor, +) -> torch.Tensor: + outputs = [] + start = 0 + for expert, count in enumerate(counts): + expert_input = inputs[start:start + count] + gate_up = F.linear(expert_input, gate_up_proj[expert]) + gate, up = gate_up.chunk(2, dim=-1) + outputs.append(F.linear(F.silu(gate) * up, down_proj[expert])) + start += count + return torch.cat(outputs, dim=0) + + +def gmm_forward( + inputs: torch.Tensor, + counts: torch.Tensor, + gate_up_weight: torch.Tensor, + down_weight: torch.Tensor, +) -> torch.Tensor: + import torch_npu + + gate_up = GmmFunction.apply(inputs, counts, gate_up_weight) + activated = torch_npu.npu_swiglu(gate_up, dim=-1) + return GmmFunction.apply(activated, counts, down_weight) + + +def main() -> None: + args = parse_args() + try: + import torch_npu # noqa: F401 + except ImportError as exc: + raise SystemExit('torch_npu is required; run this script in the Ascend container.') from exc + + if not torch.npu.is_available(): + raise SystemExit('torch.npu.is_available() is False') + + device = torch.device(args.device) + dtype = getattr(torch, args.dtype) + torch.npu.set_device(device.index or 0) + torch.manual_seed(20260901) + torch.npu.manual_seed_all(20260901) + + # Preserve the DeepSeek-V4 relation hidden == 2 * intermediate. The gate/up + # matrix is deliberately non-symmetric so an omitted transpose is visible. + experts = 2 + hidden = 64 + intermediate = 32 + token_counts = [8, 8] + inputs = torch.randn(sum(token_counts), hidden, device=device, dtype=dtype) * 0.1 + gate_up_proj = torch.randn(experts, 2 * intermediate, hidden, device=device, dtype=dtype) * 0.02 + down_proj = torch.randn(experts, hidden, intermediate, device=device, dtype=dtype) * 0.02 + module = PackedExperts(gate_up_proj, down_proj).to(device) + + normalized_gate_up, normalized_down = _normalize_packed_expert_weights(module, dtype, hidden) + counts = torch.tensor(token_counts, device=device, dtype=torch.int64) + + with torch.no_grad(): + expected = reference_forward(inputs, token_counts, gate_up_proj, down_proj) + actual = gmm_forward(inputs, counts, normalized_gate_up, normalized_down) + # Reproduce the old DeepSeek-V4 bug: the square gate/up tensor was not transposed. + old_bug = gmm_forward(inputs, counts, gate_up_proj, down_proj.transpose(1, 2)) + torch.npu.synchronize() + + difference = (actual.float() - expected.float()).abs() + old_difference = (old_bug.float() - expected.float()).abs() + passed = torch.allclose(actual.float(), expected.float(), rtol=args.rtol, atol=args.atol) + report = { + 'device': str(device), + 'dtype': str(dtype), + 'input_shape': list(inputs.shape), + 'gate_up_shape_transformers': list(gate_up_proj.shape), + 'down_shape_transformers': list(down_proj.shape), + 'gate_up_shape_gmm': list(normalized_gate_up.shape), + 'down_shape_gmm': list(normalized_down.shape), + 'rtol': args.rtol, + 'atol': args.atol, + 'max_abs_diff': difference.max().item(), + 'mean_abs_diff': difference.mean().item(), + 'old_bug_max_abs_diff': old_difference.max().item(), + 'old_bug_mean_abs_diff': old_difference.mean().item(), + 'passed': bool(passed), + } + + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2), encoding='utf-8') + print(json.dumps(report, indent=2)) + print(f'Report saved to: {args.output.resolve()}') + if not passed: + raise SystemExit(1) + + +if __name__ == '__main__': + main() diff --git a/cookbook/client/server/transformer/read_me.md b/cookbook/client/server/transformer/read_me.md index 21e725919..260279e8b 100644 --- a/cookbook/client/server/transformer/read_me.md +++ b/cookbook/client/server/transformer/read_me.md @@ -1,47 +1,546 @@ -已新增独立的“两台机器、每台 2 张昇腾 A3”配置，原来的 16 卡文件没有修改。 +# DeepSeek-V4 四层模型：FSDP/EP/NPU GMM 排查手册 -文件： +本文档用于定位 Twinkle 加载 DeepSeek-V4 后基座生成异常的问题。整套检查只需要从正式 BF16 模型截取的前四层，不依赖模型能够正常回答自然语言。 -- [4 卡服务端配置](/Users/linjiajia/project/twinkle/cookbook/client/server/transformer/server_config_dsv4_0731_npu_2node_2npu.yaml) -- [4 卡服务端启动脚本](/Users/linjiajia/project/twinkle/cookbook/client/server/transformer/run_dsv4_0731_npu_2node_2npu.sh) -- [4 卡客户端脚本](/Users/linjiajia/project/twinkle/cookbook/client/twinkle/run_dsv4_0731_npu_2node_2npu_client.sh) +排查目标是区分： -主节点 `172.61.10.111`： +1. checkpoint 或 Transformers 转换错误； +2. 两个节点加载出的源权重不一致； +3. FSDP node-local 权重分发错误； +4. EP 专家切分或 AllToAll 错误； +5. NPU grouped-matmul 专家权重方向错误。 + +四层模型无法形成有意义的自然语言回答是正常现象。所有结论必须依据权重诊断日志和最后一个位置的 logits，不能依据生成文本是否可读。 + +## 本次修改包含的文件 + +核心修复及诊断开关： + +- `src/twinkle/kernel/ops/moe/npu.py` + - 同时根据 `gate_up_proj` 和 `down_proj` 判断 Transformers `[E,out,in]` 与 GMM `[E,in,out]` 布局。 + - 修复 DeepSeek-V4 中 `hidden_size == 2 * moe_intermediate_size` 导致方阵 `gate_up_proj` 被错误识别的问题。 +- `src/twinkle/kernel/ops/ep/__init__.py` + - 新增 `TWINKLE_EP_FORCE_LOOP=1`，可强制使用逐专家 `F.linear` 参考实现。 +- `src/twinkle/model/transformers/moe/expert_parallel.py` + - 新增首轮 EP 路由、split、专家区间及输出诊断日志。 +- `src/twinkle/model/transformers/strategy/native_fsdp.py` + - 新增 node-local 源权重和 EP 本地专家切片诊断日志。 + +四层诊断配置及启动脚本： + +- `cookbook/client/server/transformer/server_config_dsv4_4layer_diag_no_ep.yaml` +- `cookbook/client/server/transformer/server_config_dsv4_4layer_diag_ep.yaml` +- `cookbook/client/server/transformer/run_dsv4_4layer_ep_diagnostic.sh` +- `cookbook/client/server/transformer/run_dsv4_0731_npu_2node_2npu.sh` + - 现在支持通过 `TWINKLE_SERVER_CONFIG_PATH` 选择配置。 + +诊断工具： + +- `cookbook/client/server/transformer/diagnostics/test_dsv4_npu_gmm_layout.py` +- `cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits.py` +- `cookbook/client/server/transformer/diagnostics/compare_dsv4_4layer_logits.py` + +单元测试： + +- `tests/kernel/ops/test_moe.py` + +## 1. 四层模型要求 + +必须使用从正式 DeepSeek-V4-Flash BF16 checkpoint 截取的前四层，不能使用随机初始化模型。配置至少需要保持： + +```text +num_hidden_layers=4 +hidden_size=4096 +moe_intermediate_size=2048 +n_routed_experts=256 +num_experts_per_tok=6 +dtype=bfloat16 +``` + +重点是保留： + +```text +hidden_size == 2 * moe_intermediate_size == 4096 +``` + +这正是原实现中方阵 `gate_up_proj` 布局误判的触发条件。 + +两个节点必须看到相同的绝对模型路径，并使用相同的 Twinkle、Transformers、PyTorch、torch-npu 和 CANN 版本。 + +启动前在两个节点分别检查： + +```bash +PROJECT_DIR=/opt/twinkle +MODEL_DIR=/highcode/shared_data/DeepSeek-V4-Flash-0731-BF16-4layers + +cd "$PROJECT_DIR" +git rev-parse HEAD +test -f "$MODEL_DIR/config.json" +test -f "$MODEL_DIR/model.safetensors.index.json" +python3 - <<'PY' +import torch +import transformers +import torch_npu + +print('torch=', torch.__version__) +print('torch_npu=', torch_npu.__version__) +print('transformers=', transformers.__version__) +PY +``` + +两个节点的输出必须一致。 + +## 2. 避免占用 Pod ephemeral-storage + +Ray 的 Unix socket 路径必须短，但缓存又不应写入只有 10 GiB 的容器临时盘。已知容器的 `/dev/shm` 有 800 GiB 时，建议把 Ray 临时目录放到共享内存；日志和结果仍写入存储卷： ```bash -DSV4_MODEL_ID=/你的/减层模型路径 \ -NETWORK_IFACE=eth0 \ -RESET_RAY=1 \ -bash cookbook/client/server/transformer/run_dsv4_0731_npu_2node_2npu.sh head +mkdir -p /dev/shm/rh +mkdir -p /highcode/shared_data/dsv4_ep_diag/tmp +mkdir -p /highcode/shared_data/dsv4_ep_diag/logs +mkdir -p /highcode/shared_data/dsv4_ep_diag/results +``` + +之后使用： + +```bash +export RAY_TMPDIR=/dev/shm/rh +export TMPDIR=/highcode/shared_data/dsv4_ep_diag/tmp +export RAY_ROTATION_MAX_BYTES=20971520 +export RAY_ROTATION_BACKUP_COUNT=1 +``` + +不要把 Ray 临时目录设成很长的路径，否则可能再次触发 AF_UNIX 107 字节路径上限。 + +## 3. 先做单卡 NPU GMM 数值测试 + +该步骤不启动 Ray，也不加载四层模型。它使用保持 DeepSeek-V4 比例的微型专家，直接比较： + +```text +F.linear 参考结果 +vs +npu_grouped_matmul 结果 ``` -Worker 节点 `172.61.12.165`： +在一台 NPU 机器执行： ```bash -DSV4_MODEL_ID=/你的/减层模型路径 \ -NETWORK_IFACE=eth0 \ -RESET_RAY=1 \ -bash cookbook/client/server/transformer/run_dsv4_0731_npu_2node_2npu.sh worker +cd /opt/twinkle + +ASCEND_RT_VISIBLE_DEVICES=0 \ +PYTHONPATH=/opt/twinkle/src \ +python3 cookbook/client/server/transformer/diagnostics/test_dsv4_npu_gmm_layout.py \ + --device npu:0 \ + --output /highcode/shared_data/dsv4_ep_diag/results/npu_gmm_layout.json +``` + +成功时： + +```text +passed=true +max_abs_diff 和 mean_abs_diff 在 BF16 允许范围内 +old_bug_max_abs_diff 明显大于修复后的 max_abs_diff +``` + +如果这里失败，不要启动分布式服务，先处理 NPU GMM 与 `F.linear` 的数值差异。 + +## 4. 三组服务端对照模式 + +三组模式只改变 EP 和专家计算方式： + +| 模式 | FSDP | EP | 专家计算 | +|---|---:|---:|---| +| `no_ep` | 开启 | 关闭 | Transformers 原始前向 | +| `ep_loop` | 开启 | 开启 | 强制逐专家 `F.linear` | +| `ep_gmm` | 开启 | 开启 | NPU grouped-matmul | + +两份 YAML 除 EP 设置外保持一致： + +```text +两节点 +每节点 2 张 NPU +world_size=4 +fsdp_size=4 +ep_size=4（仅 EP 模式） +memory_efficient_init=true +max_loras=1 +max_length=512 ``` -客户端： +EP=4 时，256 个专家的预期区间为： + +```text +rank 0 / ep_rank 0: experts 0..63 +rank 1 / ep_rank 1: experts 64..127 +rank 2 / ep_rank 2: experts 128..191 +rank 3 / ep_rank 3: experts 192..255 +``` + +## 5. 两个节点的公共环境 + +以下示例使用： + +```text +Head: 172.61.10.111 +Worker: 172.61.12.165 +网卡: eth0 +``` + +如果实际环境使用 `bond0` 或其他地址，只修改环境变量，不需要修改脚本或 YAML。 + +Head 节点执行： + +```bash +cd /opt/twinkle + +export DSV4_MODEL_ID=/highcode/shared_data/DeepSeek-V4-Flash-0731-BF16-4layers +export HEAD_IP=172.61.10.111 +export WORKER_IP=172.61.12.165 +export NODE_IP=172.61.10.111 +export NETWORK_IFACE=eth0 +export ASCEND_RT_VISIBLE_DEVICES=0,1 +export RESET_RAY=1 +export RAY_TMPDIR=/dev/shm/rh +export TMPDIR=/highcode/shared_data/dsv4_ep_diag/tmp +export RAY_ROTATION_MAX_BYTES=20971520 +export RAY_ROTATION_BACKUP_COUNT=1 +``` + +Worker 节点执行： + +```bash +cd /opt/twinkle + +export DSV4_MODEL_ID=/highcode/shared_data/DeepSeek-V4-Flash-0731-BF16-4layers +export HEAD_IP=172.61.10.111 +export WORKER_IP=172.61.12.165 +export NODE_IP=172.61.12.165 +export NETWORK_IFACE=eth0 +export ASCEND_RT_VISIBLE_DEVICES=0,1 +export RESET_RAY=1 +export RAY_TMPDIR=/dev/shm/rh +export TMPDIR=/highcode/shared_data/dsv4_ep_diag/tmp +export RAY_ROTATION_MAX_BYTES=20971520 +export RAY_ROTATION_BACKUP_COUNT=1 +``` + +必须确保两个节点选择相同模式。 + +## 6. 模式一：启动 no-EP 基准 + +先在 Head 节点启动；脚本会等待 Worker 加入： + +```bash +nohup bash cookbook/client/server/transformer/run_dsv4_4layer_ep_diagnostic.sh head no_ep \ + >/highcode/shared_data/dsv4_ep_diag/logs/no_ep_head.log 2>&1 /highcode/shared_data/dsv4_ep_diag/logs/no_ep_worker.log 2>&1 /highcode/shared_data/dsv4_ep_diag/logs/ep_loop_head.log 2>&1 /highcode/shared_data/dsv4_ep_diag/logs/ep_loop_worker.log 2>&1 /highcode/shared_data/dsv4_ep_diag/logs/ep_gmm_head.log 2>&1 /highcode/shared_data/dsv4_ep_diag/logs/ep_gmm_worker.log 2>&1 /highcode/shared_data/dsv4_ep_diag/results/ep_diagnostics.txt +``` + +重点检查以下日志。 + +### 11.1 两个节点的源权重 + +每个 node-local source rank 都会输出： + +```text +local_source= +param=<参数名> +full_shape=<完整专家形状> +source_preview=<采样值> +``` + +相同参数在两个 node-local source 上必须满足： + +```text +full_shape 相同 +source_preview 相同 +``` + +如果不同，优先检查两个节点： + +- 是否使用相同模型目录； +- 模型文件是否一致； +- Transformers/torch版本是否一致； +- 是否实际运行同一个 Twinkle commit。 + +### 11.2 EP 专家范围 + +每个rank会输出： + +```text +rank= +ep_rank= +expert_range=[start,end) +local_shape=<本地权重形状> +local_preview=<本地切片采样值> +``` + +EP=4时必须严格对应： + +```text +rank0 -> [0,64) +rank1 -> [64,128) +rank2 -> [128,192) +rank3 -> [192,256) +``` + +本地权重预期形状： + +```text +gate_up_proj: [64,4096,4096] +down_proj: [64,4096,2048] +``` + +### 11.3 Router与AllToAll split + +每层第一次前向会输出： + +```text +selected_experts +routing_weights +input_splits +output_splits +output_finite +``` + +固定输入下，各rank在AllToAll前的 `selected_experts` 和 `routing_weights` 应一致。代码还会验证： + +```text +sum(input_splits) == token_count * num_experts_per_tok +``` + +如果该条件不成立，诊断模式会直接报错退出。 + +## 12. 单元测试和静态检查 + +在安装了项目依赖的环境执行： + +```bash +cd /opt/twinkle + +python3 -m pytest -q tests/kernel/ops/test_moe.py + +python3 -m py_compile \ + cookbook/client/server/transformer/diagnostics/test_dsv4_npu_gmm_layout.py \ + cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits.py \ + cookbook/client/server/transformer/diagnostics/compare_dsv4_4layer_logits.py + +bash -n \ + cookbook/client/server/transformer/run_dsv4_0731_npu_2node_2npu.sh \ + cookbook/client/server/transformer/run_dsv4_4layer_ep_diagnostic.sh +``` + +## 13. 完整模型的最终验证 + +只有满足以下条件后才启动43层完整模型： + +```text +单卡 NPU GMM 与 F.linear 一致 +no_ep logits ≈ ep_loop logits +ep_loop logits ≈ ep_gmm logits +两个 node-local source 的权重采样一致 +EP 专家范围正确 +路由 split 不变量通过 +``` + +完整模型启动时应关闭诊断开关： + +```bash +export TWINKLE_EP_FORCE_LOOP=0 +export TWINKLE_EP_DIAGNOSTICS=0 +``` + +然后按顺序验证： + +1. 不训练、不注册旧LoRA，先检查基座生成； +2. 注册全零LoRA，确认它与 `disable_lora=True` 逐token一致； +3. 修复前在错误EP前向下训练的LoRA不能作为正确性基准； +4. 基座确认正常后重新训练LoRA； +5. 对比重新训练后的 base/logprob/生成结果。 + +## 14. 清理 + +诊断结束后在两个节点执行： ```bash -DSV4_MODEL_ID=/你的/减层模型路径 \ -DATASET_ID=/model/ljl/dataset/self-cognition.jsonl \ -OUTPUT_DIR=/shared/twinkle_output/dsv4-0731-a3-2node-4npu \ -bash cookbook/client/twinkle/run_dsv4_0731_npu_2node_2npu_client.sh +ray stop --force ``` -两个 IP 已写成默认值，通常不需要再传 `HEAD_IP` 或 `NODE_IP`。配置为： +结果位于共享卷： ```text -每节点 NPU：2 -节点数：2 -总 rank：4 -FSDP size：4 -EP size：4 -默认 batch size：4 +/highcode/shared_data/dsv4_ep_diag/results +/highcode/shared_data/dsv4_ep_diag/logs ``` -`DSV4_MODEL_ID` 已直接接入 YAML 环境变量解析，因此替换模型时不需要修改 YAML，只需确保主节点、Worker 和客户端传入相同值。语法和配置拓扑校验均已通过。 +确认不再需要后可手动清理共享卷中的诊断日志和结果。`/dev/shm/rh` 仅用于当前 Pod 的 Ray 临时文件，Pod 删除后不会保留。 diff --git a/cookbook/client/server/transformer/run_dsv4_0731_npu_2node_2npu.sh b/cookbook/client/server/transformer/run_dsv4_0731_npu_2node_2npu.sh index 7f3c9404c..814a9eb94 100755 --- a/cookbook/client/server/transformer/run_dsv4_0731_npu_2node_2npu.sh +++ b/cookbook/client/server/transformer/run_dsv4_0731_npu_2node_2npu.sh @@ -31,7 +31,10 @@ fi SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" PROJECT_DIR="$(cd -- "$SCRIPT_DIR/../../../.." && pwd)" -CONFIG_PATH="$SCRIPT_DIR/server_config_dsv4_0731_npu_2node_2npu.yaml" +CONFIG_PATH="${TWINKLE_SERVER_CONFIG_PATH:-$SCRIPT_DIR/server_config_dsv4_0731_npu_2node_2npu.yaml}" +if [[ "$CONFIG_PATH" != /* ]]; then + CONFIG_PATH="$PROJECT_DIR/$CONFIG_PATH" +fi export DSV4_MODEL_ID="${DSV4_MODEL_ID:-hf://deepseek-ai/DeepSeek-V4-Flash-0731}" DATASET_PATH="${DATASET_ID:-/model/ljl/dataset/self-cognition.jsonl}" @@ -55,6 +58,8 @@ export ASCEND_RT_VISIBLE_DEVICES="${ASCEND_RT_VISIBLE_DEVICES:-0,1}" export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1 export TWINKLE_TRUST_REMOTE_CODE=1 export TWINKLE_FAIL_FAST=1 +export TWINKLE_EP_FORCE_LOOP="${TWINKLE_EP_FORCE_LOOP:-0}" +export TWINKLE_EP_DIAGNOSTICS="${TWINKLE_EP_DIAGNOSTICS:-0}" export TOKENIZERS_PARALLELISM=true export GLOO_SOCKET_IFNAME="$NETWORK_IFACE" export HCCL_SOCKET_IFNAME="$NETWORK_IFACE" @@ -70,7 +75,9 @@ if [[ "${#VISIBLE_NPUS[@]}" -ne "$NPU_PER_NODE" ]]; then fi test -f "$CONFIG_PATH" -test -f "$DATASET_PATH" +if [[ "${SKIP_DATASET_CHECK:-0}" != "1" ]]; then + test -f "$DATASET_PATH" +fi if [[ "$DSV4_MODEL_ID" == hf://* || "$DSV4_MODEL_ID" == ms://* ]]; then echo "Model will be downloaded through the configured Hub backend: $DSV4_MODEL_ID" @@ -157,5 +164,7 @@ finally: PY python3 -m twinkle.server check-config --config "$CONFIG_PATH" +echo "Twinkle config: $CONFIG_PATH" +echo "EP diagnostics: force_loop=$TWINKLE_EP_FORCE_LOOP diagnostics=$TWINKLE_EP_DIAGNOSTICS" echo "Launching Twinkle Server at http://$HEAD_IP:8000" exec python3 -m twinkle.server launch --config "$CONFIG_PATH" diff --git a/cookbook/client/server/transformer/run_dsv4_4layer_ep_diagnostic.sh b/cookbook/client/server/transformer/run_dsv4_4layer_ep_diagnostic.sh new file mode 100755 index 000000000..e1821f91f --- /dev/null +++ b/cookbook/client/server/transformer/run_dsv4_4layer_ep_diagnostic.sh @@ -0,0 +1,37 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROLE="${1:-}" +MODE="${2:-}" + +if [[ "$ROLE" != "head" && "$ROLE" != "worker" ]]; then + echo "Usage: $0 {head|worker} {no_ep|ep_loop|ep_gmm}" >&2 + exit 2 +fi +if [[ "$MODE" != "no_ep" && "$MODE" != "ep_loop" && "$MODE" != "ep_gmm" ]]; then + echo "Usage: $0 {head|worker} {no_ep|ep_loop|ep_gmm}" >&2 + exit 2 +fi + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" + +case "$MODE" in + no_ep) + export TWINKLE_SERVER_CONFIG_PATH="$SCRIPT_DIR/server_config_dsv4_4layer_diag_no_ep.yaml" + export TWINKLE_EP_FORCE_LOOP=0 + ;; + ep_loop) + export TWINKLE_SERVER_CONFIG_PATH="$SCRIPT_DIR/server_config_dsv4_4layer_diag_ep.yaml" + export TWINKLE_EP_FORCE_LOOP=1 + ;; + ep_gmm) + export TWINKLE_SERVER_CONFIG_PATH="$SCRIPT_DIR/server_config_dsv4_4layer_diag_ep.yaml" + export TWINKLE_EP_FORCE_LOOP=0 + ;; +esac + +export TWINKLE_EP_DIAGNOSTICS="${TWINKLE_EP_DIAGNOSTICS:-1}" +export SKIP_DATASET_CHECK=1 + +echo "Starting DeepSeek-V4 four-layer diagnostic: role=$ROLE mode=$MODE" +exec "$SCRIPT_DIR/run_dsv4_0731_npu_2node_2npu.sh" "$ROLE" diff --git a/cookbook/client/server/transformer/server_config_dsv4_4layer_diag_ep.yaml b/cookbook/client/server/transformer/server_config_dsv4_4layer_diag_ep.yaml new file mode 100644 index 000000000..13d9dabe6 --- /dev/null +++ b/cookbook/client/server/transformer/server_config_dsv4_4layer_diag_ep.yaml @@ -0,0 +1,117 @@ +# Four-layer DeepSeek-V4 diagnostic: FSDP2 + EP on two nodes, two NPUs per node. +# DSV4_MODEL_ID must point to a structurally complete four-layer BF16 checkpoint. + +proxy_location: EveryNode + +http_options: + host: 0.0.0.0 + port: 8000 + +telemetry: + enabled: false + otlp_endpoint: http://localhost:4317 + +persistence: + mode: memory + key_prefix: dsv4-4layer-diag-ep + +applications: + - name: server + route_prefix: /api/v1 + import_path: server + args: + server_config: + per_token_model_limit: 1 + supported_models: + - deepseek-v4-0731-local + deployments: + - name: TinkerCompatServer + max_ongoing_requests: 4 + autoscaling_config: + min_replicas: 1 + max_replicas: 1 + target_ongoing_requests: 1 + ray_actor_options: + num_cpus: 0.1 + runtime_env: + env_vars: + TWINKLE_FAIL_FAST: "1" + + - name: models-deepseek-v4-0731-local + route_prefix: /api/v1/model/deepseek-v4-0731-local + import_path: model + args: + backend: transformers + model_id: ${oc.env:DSV4_MODEL_ID} + nproc_per_node: 2 + max_loras: 1 + max_r: 8 + max_length: 512 + mixed_precision: bf16 + strategy: native_fsdp + memory_efficient_init: true + target_modules: all-linear + fsdp_config: + reshard_after_forward: true + expert_parallel: + enabled: true + ep_size: 4 + router_dtype: fp32 + keep_router_logits: false + device_group: + name: model + ranks: 4 + device_type: npu + device_mesh: + device_type: npu + fsdp_size: 4 + dp_size: 1 + ep_size: 4 + queue_config: + rps_limit: 1 + tps_limit: 10000 + max_input_tokens: 512 + queue_timeout: 7200 + execution_timeout: 7200 + adapter_config: + adapter_timeout: 60 + adapter_max_lifetime: 7200 + deployments: + - name: ModelManagement + autoscaling_config: + min_replicas: 1 + max_replicas: 1 + target_ongoing_requests: 1 + ray_actor_options: + num_cpus: 0.1 + runtime_env: + env_vars: + TWINKLE_TRUST_REMOTE_CODE: "1" + TWINKLE_FAIL_FAST: "1" + TWINKLE_EP_FORCE_LOOP: ${oc.env:TWINKLE_EP_FORCE_LOOP,0} + TWINKLE_EP_DIAGNOSTICS: ${oc.env:TWINKLE_EP_DIAGNOSTICS,1} + + - name: processor + route_prefix: /api/v1/processor + import_path: processor + args: + ncpu_proc_per_node: 1 + device_group: + name: processor + ranks: 2 + device_type: CPU + device_mesh: + device_type: CPU + dp_size: 2 + deployments: + - name: ProcessorManagement + autoscaling_config: + min_replicas: 1 + max_replicas: 1 + target_ongoing_requests: 1 + ray_actor_options: + num_cpus: 0.1 + runtime_env: + env_vars: + TWINKLE_TRUST_REMOTE_CODE: "1" + TWINKLE_FAIL_FAST: "1" diff --git a/cookbook/client/server/transformer/server_config_dsv4_4layer_diag_no_ep.yaml b/cookbook/client/server/transformer/server_config_dsv4_4layer_diag_no_ep.yaml new file mode 100644 index 000000000..3576f7794 --- /dev/null +++ b/cookbook/client/server/transformer/server_config_dsv4_4layer_diag_no_ep.yaml @@ -0,0 +1,117 @@ +# Four-layer DeepSeek-V4 diagnostic reference: FSDP2 without expert parallel. +# Apart from the EP settings, keep this aligned with server_config_dsv4_4layer_diag_ep.yaml. + +proxy_location: EveryNode + +http_options: + host: 0.0.0.0 + port: 8000 + +telemetry: + enabled: false + otlp_endpoint: http://localhost:4317 + +persistence: + mode: memory + key_prefix: dsv4-4layer-diag-no-ep + +applications: + - name: server + route_prefix: /api/v1 + import_path: server + args: + server_config: + per_token_model_limit: 1 + supported_models: + - deepseek-v4-0731-local + deployments: + - name: TinkerCompatServer + max_ongoing_requests: 4 + autoscaling_config: + min_replicas: 1 + max_replicas: 1 + target_ongoing_requests: 1 + ray_actor_options: + num_cpus: 0.1 + runtime_env: + env_vars: + TWINKLE_FAIL_FAST: "1" + + - name: models-deepseek-v4-0731-local + route_prefix: /api/v1/model/deepseek-v4-0731-local + import_path: model + args: + backend: transformers + model_id: ${oc.env:DSV4_MODEL_ID} + nproc_per_node: 2 + max_loras: 1 + max_r: 8 + max_length: 512 + mixed_precision: bf16 + strategy: native_fsdp + memory_efficient_init: true + target_modules: all-linear + fsdp_config: + reshard_after_forward: true + expert_parallel: + enabled: false + ep_size: 1 + router_dtype: fp32 + keep_router_logits: false + device_group: + name: model + ranks: 4 + device_type: npu + device_mesh: + device_type: npu + fsdp_size: 4 + dp_size: 1 + ep_size: 1 + queue_config: + rps_limit: 1 + tps_limit: 10000 + max_input_tokens: 512 + queue_timeout: 7200 + execution_timeout: 7200 + adapter_config: + adapter_timeout: 60 + adapter_max_lifetime: 7200 + deployments: + - name: ModelManagement + autoscaling_config: + min_replicas: 1 + max_replicas: 1 + target_ongoing_requests: 1 + ray_actor_options: + num_cpus: 0.1 + runtime_env: + env_vars: + TWINKLE_TRUST_REMOTE_CODE: "1" + TWINKLE_FAIL_FAST: "1" + TWINKLE_EP_FORCE_LOOP: "0" + TWINKLE_EP_DIAGNOSTICS: ${oc.env:TWINKLE_EP_DIAGNOSTICS,1} + + - name: processor + route_prefix: /api/v1/processor + import_path: processor + args: + ncpu_proc_per_node: 1 + device_group: + name: processor + ranks: 2 + device_type: CPU + device_mesh: + device_type: CPU + dp_size: 2 + deployments: + - name: ProcessorManagement + autoscaling_config: + min_replicas: 1 + max_replicas: 1 + target_ongoing_requests: 1 + ray_actor_options: + num_cpus: 0.1 + runtime_env: + env_vars: + TWINKLE_TRUST_REMOTE_CODE: "1" + TWINKLE_FAIL_FAST: "1" diff --git a/src/twinkle/kernel/ops/ep/__init__.py b/src/twinkle/kernel/ops/ep/__init__.py index c011d3ced..556d1cff7 100644 --- a/src/twinkle/kernel/ops/ep/__init__.py +++ b/src/twinkle/kernel/ops/ep/__init__.py @@ -11,6 +11,7 @@ """ from __future__ import annotations +import os import torch from abc import ABC, abstractmethod from torch import nn @@ -51,6 +52,11 @@ def forward( _IMPLS: list[EpExpertsGmm] | None = None _PATH_LOGGED = False _WARN_LOGGED = False +_FORCE_LOOP_LOGGED = False + + +def _env_flag(name: str) -> bool: + return os.environ.get(name, '').strip().lower() in {'1', 'true', 'yes', 'on'} def _get_impls() -> list[EpExpertsGmm]: @@ -59,6 +65,14 @@ def _get_impls() -> list[EpExpertsGmm]: Each backend module is imported defensively: platforms lacking its dependencies (e.g. no torch_npu) simply skip that backend. """ + global _FORCE_LOOP_LOGGED + if _env_flag('TWINKLE_EP_FORCE_LOOP'): + from .loop import LoopEpExpertsGmm + if not _FORCE_LOOP_LOGGED: + logger.warning('EP experts compute: TWINKLE_EP_FORCE_LOOP=1; forcing the per-expert F.linear loop.') + _FORCE_LOOP_LOGGED = True + return [LoopEpExpertsGmm()] + global _IMPLS if _IMPLS is None: _IMPLS = [] diff --git a/src/twinkle/kernel/ops/moe/npu.py b/src/twinkle/kernel/ops/moe/npu.py index 1f847669b..8fcbfa694 100644 --- a/src/twinkle/kernel/ops/moe/npu.py +++ b/src/twinkle/kernel/ops/moe/npu.py @@ -62,18 +62,32 @@ def npu_grouped_mm(input: torch.Tensor, weight_ekn: torch.Tensor, offs: torch.Te def _normalize_packed_expert_weights(module, input_dtype, hidden_dim): gate_up_proj = module.gate_up_proj.to(input_dtype) down_proj = module.down_proj.to(input_dtype) - if gate_up_proj.shape[1] == hidden_dim: - gate_up_weight = gate_up_proj - elif gate_up_proj.shape[2] == hidden_dim: + if gate_up_proj.ndim != 3 or down_proj.ndim != 3: + raise RuntimeError( + 'Packed expert weights must be 3D: ' + f'gate_up_proj={tuple(gate_up_proj.shape)}, down_proj={tuple(down_proj.shape)}.') + + # torch.nn.functional.linear stores weights as [out_features, in_features], + # while torch_npu.npu_grouped_matmul consumes [in_features, out_features]. + # Infer one layout from *both* tensors instead of inspecting gate_up_proj in + # isolation. DeepSeek-V4 has hidden_size == 2 * intermediate_size, so its + # gate_up_proj is square and either of its last two dimensions can look like + # the input dimension. down_proj remains non-square and disambiguates it. + linear_layout = gate_up_proj.shape[2] == hidden_dim and down_proj.shape[1] == hidden_dim + grouped_mm_layout = gate_up_proj.shape[1] == hidden_dim and down_proj.shape[2] == hidden_dim + if linear_layout == grouped_mm_layout: + raise RuntimeError( + 'Unable to determine packed expert weight layout: ' + f'gate_up_proj={tuple(gate_up_proj.shape)}, down_proj={tuple(down_proj.shape)}, ' + f'hidden_dim={hidden_dim}. Expected either Transformers/F.linear ' + '[E, out, in] tensors or grouped-matmul [E, in, out] tensors.') + + if linear_layout: gate_up_weight = gate_up_proj.transpose(1, 2) - else: - raise RuntimeError(f'Unsupported gate_up_proj shape: {tuple(gate_up_proj.shape)}.') - if down_proj.shape[2] == hidden_dim: - down_weight = down_proj - elif down_proj.shape[1] == hidden_dim: down_weight = down_proj.transpose(1, 2) else: - raise RuntimeError(f'Unsupported down_proj shape: {tuple(down_proj.shape)}.') + gate_up_weight = gate_up_proj + down_weight = down_proj return gate_up_weight, down_weight diff --git a/src/twinkle/model/transformers/moe/expert_parallel.py b/src/twinkle/model/transformers/moe/expert_parallel.py index aa828d24b..0bc3a391d 100644 --- a/src/twinkle/model/transformers/moe/expert_parallel.py +++ b/src/twinkle/model/transformers/moe/expert_parallel.py @@ -2,16 +2,24 @@ from __future__ import annotations import inspect +import os import torch import torch.distributed as dist from dataclasses import dataclass from torch import nn from typing import Any, Dict, Iterable, List, Optional, Tuple +from twinkle import get_logger from twinkle.kernel.ops import ep_forward from twinkle.model.transformers.moe.ep_utils import preprocess, token_pre_all2all, tokens_post_all2all from twinkle.utils import DeviceMesh +logger = get_logger() + + +def _ep_diagnostics_enabled() -> bool: + return os.environ.get('TWINKLE_EP_DIAGNOSTICS', '').strip().lower() in {'1', 'true', 'yes', 'on'} + @dataclass class ExpertParallelConfig: @@ -262,6 +270,32 @@ def forward(hidden_states: torch.Tensor, *args, **kwargs): num_global_sum_tokens_per_local_expert, ) = preprocess(expert_mask, num_experts, ep_group) + log_diagnostics = _ep_diagnostics_enabled() and not getattr(block, '_ep_diagnostics_logged', False) + if log_diagnostics: + expected_assignments = hidden_states_2d.shape[0] * top_k + actual_assignments = int(input_splits.sum().item()) + if actual_assignments != expected_assignments: + raise RuntimeError( + f'EP routing assignment mismatch for {block_name}: ' + f'input_splits={actual_assignments}, expected={expected_assignments}.') + selected_preview = selected_experts[:min(8, selected_experts.shape[0])].detach().cpu().tolist() + weights_preview = routing_weights[:min(8, routing_weights.shape[0])].float().detach().cpu().tolist() + logger.warning( + '[EP_DIAG] rank=%s ep_rank=%s block=%s expert_range=[%s,%s) tokens=%s top_k=%s ' + 'input_splits=%s output_splits=%s selected_experts=%s routing_weights=%s', + dist.get_rank(), + block._ep_rank, + block_name, + block._ep_local_start, + block._ep_local_end, + hidden_states_2d.shape[0], + top_k, + input_splits.tolist(), + output_splits.tolist(), + selected_preview, + weights_preview, + ) + # 2. token_pre_all2all: permute → all_to_all → sort_chunks ( global_permuted_hidden_states, @@ -317,6 +351,21 @@ def forward(hidden_states: torch.Tensor, *args, **kwargs): if len(orig_shape) == 3: final_hidden = final_hidden.view(batch_size, seq_len, hidden_dim) + if log_diagnostics: + flat_output = final_hidden.detach().reshape(-1) + preview = flat_output[:min(8, flat_output.numel())].float().cpu().tolist() + finite = bool(torch.isfinite(final_hidden).all().item()) + logger.warning( + '[EP_DIAG] rank=%s ep_rank=%s block=%s output_shape=%s output_finite=%s output_preview=%s', + dist.get_rank(), + block._ep_rank, + block_name, + tuple(final_hidden.shape), + finite, + preview, + ) + block._ep_diagnostics_logged = True + if cfg.keep_router_logits and returns_router_logits: return final_hidden, router_logits return final_hidden diff --git a/src/twinkle/model/transformers/strategy/native_fsdp.py b/src/twinkle/model/transformers/strategy/native_fsdp.py index 2671284e7..f1bd6c6ec 100644 --- a/src/twinkle/model/transformers/strategy/native_fsdp.py +++ b/src/twinkle/model/transformers/strategy/native_fsdp.py @@ -1000,6 +1000,9 @@ def _broadcast_sharded_state_dict( rank_to_ep_rank = rank_to_ep_rank or {} adapter_source_sd = adapter_source_sd or {} adapter_full_sd = adapter_full_sd or {} + diagnostics_enabled = os.environ.get('TWINKLE_EP_DIAGNOSTICS', '').strip().lower() in { + '1', 'true', 'yes', 'on' + } source_metadata = None source_keys = None adapter_metadata = None @@ -1113,14 +1116,43 @@ def _scatter_ep_expert_tensor(param_name, full_tensor, sharded_param): raise RuntimeError(f"EP expert parameter '{param_name}' expects {num_experts} experts, " f'but source state has shape {tuple(full_tensor.shape)}. ' 'Rank0 must capture the full pre-EP state_dict before apply_expert_parallel().') + if diagnostics_enabled: + logger.warning( + '[EP_DIAG] rank=%s local_source=%s param=%s full_shape=%s source_preview=%s', + rank, + local_source_rank, + param_name, + tuple(full_tensor.shape), + _diagnostic_tensor_preview(full_tensor), + ) local_tensor = _scatter_ep_tensor_from_source( full_tensor, local_tensor, shard_dim=0, shard_size=experts_per_rank, ) + if diagnostics_enabled: + ep_rank = rank_to_ep_rank[rank] + start = ep_rank * experts_per_rank + logger.warning( + '[EP_DIAG] rank=%s ep_rank=%s param=%s expert_range=[%s,%s) local_shape=%s local_preview=%s', + rank, + ep_rank, + param_name, + start, + start + experts_per_rank, + tuple(local_tensor.shape), + _diagnostic_tensor_preview(local_tensor), + ) return local_tensor + def _diagnostic_tensor_preview(tensor: torch.Tensor) -> List[float]: + flat = tensor.detach().reshape(-1) + if flat.numel() == 0: + return [] + indices = sorted({0, flat.numel() // 3, (2 * flat.numel()) // 3, flat.numel() - 1}) + return flat[indices].float().cpu().tolist() + def _scatter_ep_tensor_from_source(full_tensor, local_tensor, *, shard_dim: int, shard_size: int): if is_source_rank: if full_tensor is None: diff --git a/tests/kernel/ops/test_moe.py b/tests/kernel/ops/test_moe.py index 4b210029c..e77bc9b6a 100644 --- a/tests/kernel/ops/test_moe.py +++ b/tests/kernel/ops/test_moe.py @@ -1,3 +1,10 @@ +import importlib + +import pytest +import torch +from torch import nn + + def test_moe_imports(): from twinkle.kernel.ops.moe.npu import ( GmmFunction, @@ -5,8 +12,64 @@ def test_moe_imports(): npu_packed_moe_experts_forward, npu_qwen3_5_moe_sparse_block_forward, ) - import torch assert issubclass(GmmFunction, torch.autograd.Function) assert callable(npu_grouped_mm) assert callable(npu_packed_moe_experts_forward) - assert callable(npu_qwen3_5_moe_sparse_block_forward) \ No newline at end of file + assert callable(npu_qwen3_5_moe_sparse_block_forward) + + +class _PackedExperts(nn.Module): + + def __init__(self, gate_up_proj: torch.Tensor, down_proj: torch.Tensor): + super().__init__() + self.gate_up_proj = nn.Parameter(gate_up_proj) + self.down_proj = nn.Parameter(down_proj) + + +def test_normalize_packed_expert_weights_resolves_square_dsv4_gate_from_down_proj(): + from twinkle.kernel.ops.moe.npu import _normalize_packed_expert_weights + + # DeepSeek-V4 relation: hidden == 2 * intermediate. gate_up_proj is + # therefore square even though it still uses F.linear [out, in] layout. + experts, hidden, intermediate = 2, 8, 4 + gate_up = torch.arange(experts * hidden * hidden, dtype=torch.float32).reshape(experts, hidden, hidden) + down = torch.arange(experts * hidden * intermediate, + dtype=torch.float32).reshape(experts, hidden, intermediate) + module = _PackedExperts(gate_up, down) + + normalized_gate_up, normalized_down = _normalize_packed_expert_weights(module, torch.float32, hidden) + + assert torch.equal(normalized_gate_up, gate_up.transpose(1, 2)) + assert torch.equal(normalized_down, down.transpose(1, 2)) + + +def test_normalize_packed_expert_weights_keeps_grouped_mm_layout(): + from twinkle.kernel.ops.moe.npu import _normalize_packed_expert_weights + + experts, hidden, intermediate = 2, 8, 4 + gate_up = torch.randn(experts, hidden, intermediate * 2) + down = torch.randn(experts, intermediate, hidden) + module = _PackedExperts(gate_up, down) + + normalized_gate_up, normalized_down = _normalize_packed_expert_weights(module, torch.float32, hidden) + + assert torch.equal(normalized_gate_up, gate_up) + assert torch.equal(normalized_down, down) + + +def test_normalize_packed_expert_weights_rejects_inconsistent_layout(): + from twinkle.kernel.ops.moe.npu import _normalize_packed_expert_weights + + module = _PackedExperts(torch.randn(2, 7, 9), torch.randn(2, 5, 6)) + with pytest.raises(RuntimeError, match='Unable to determine packed expert weight layout'): + _normalize_packed_expert_weights(module, torch.float32, hidden_dim=8) + + +def test_ep_force_loop_environment_switch(monkeypatch): + ep_ops = importlib.import_module('twinkle.kernel.ops.ep') + + monkeypatch.setenv('TWINKLE_EP_FORCE_LOOP', '1') + implementations = ep_ops._get_impls() + + assert len(implementations) == 1 + assert implementations[0].name == 'per-expert loop' From e329ee6bae7f01b631e0e8be3bceeede05e800f8 Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Wed, 2 Sep 2026 10:19:12 +0800 Subject: [PATCH 20/28] wip --- .../diagnostics/probe_dsv4_4layer_logits.py | 53 +++++++++++++++++-- 1 file changed, 50 insertions(+), 3 deletions(-) diff --git a/cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits.py b/cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits.py index dd40d122d..ea9d9ee7a 100755 --- a/cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits.py +++ b/cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits.py @@ -6,6 +6,7 @@ import argparse import hashlib import json +import time from pathlib import Path from typing import Any @@ -28,6 +29,18 @@ def parse_args() -> argparse.Namespace: parser.add_argument('--served-model', default='deepseek-v4-0731-local') parser.add_argument('--output-dir', type=Path, default=Path('output/dsv4_ep_diag')) parser.add_argument('--input-ids', default=','.join(str(item) for item in DEFAULT_INPUT_IDS)) + parser.add_argument( + '--capacity-wait-seconds', + type=float, + default=90.0, + help='Wait this long for a previous diagnostic adapter to expire (default: 90).', + ) + parser.add_argument( + '--capacity-poll-seconds', + type=float, + default=5.0, + help='Seconds between LoRA-capacity checks (default: 5).', + ) return parser.parse_args() @@ -59,6 +72,38 @@ def _tensor_sha256(tensor: torch.Tensor) -> str: return hashlib.sha256(values).hexdigest() +def _wait_for_free_lora(client: Any, timeout: float, poll_interval: float) -> Any: + """Wait for the prior probe's session-bound adapter to expire.""" + if timeout < 0: + raise ValueError('--capacity-wait-seconds must be non-negative') + if poll_interval <= 0: + raise ValueError('--capacity-poll-seconds must be positive') + + deadline = time.monotonic() + timeout + while True: + capacity = client.get_capacity_info() + if capacity.free_loras >= 1: + return capacity + + remaining = deadline - time.monotonic() + if remaining <= 0: + raise RuntimeError( + 'Timed out waiting for a free LoRA slot: ' + f'max_loras={capacity.max_loras}, used_loras={capacity.used_loras}, ' + f'free_loras={capacity.free_loras}. The previous diagnostic adapter ' + 'did not expire; inspect the ModelManagement cleanup logs or restart ' + 'the diagnostic server.') + + sleep_seconds = min(poll_interval, remaining) + print( + 'Waiting for a free LoRA slot: ' + f'used={capacity.used_loras}/{capacity.max_loras}, ' + f'remaining={remaining:.1f}s', + flush=True, + ) + time.sleep(sleep_seconds) + + def main() -> None: args = parse_args() input_ids = [int(item.strip()) for item in args.input_ids.split(',') if item.strip()] @@ -71,9 +116,11 @@ def main() -> None: session_heartbeat_interval=10, ) try: - capacity = client.get_capacity_info() - if capacity.free_loras < 1: - raise RuntimeError('No free LoRA slot. Restart the diagnostic server before running the probe.') + _wait_for_free_lora( + client, + timeout=args.capacity_wait_seconds, + poll_interval=args.capacity_poll_seconds, + ) model = MultiLoraTransformersModel(model_id=args.served_model) model.add_adapter_to_model( From 2d89fb2557cede538a12f3dd72f76c1f224c2485 Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Wed, 2 Sep 2026 10:28:29 +0800 Subject: [PATCH 21/28] wip --- .../diagnostics/probe_dsv4_4layer_logits.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits.py b/cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits.py index ea9d9ee7a..30575c592 100755 --- a/cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits.py +++ b/cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits.py @@ -87,6 +87,12 @@ def _wait_for_free_lora(client: Any, timeout: float, poll_interval: float) -> An remaining = deadline - time.monotonic() if remaining <= 0: + if capacity.max_loras == 0: + raise RuntimeError( + 'Timed out waiting for ModelManagement capacity: ' + 'max_loras=0, used_loras=0. No model replica registered its ' + 'capacity; inspect the ModelManagement startup and replica ' + 'registration logs.') raise RuntimeError( 'Timed out waiting for a free LoRA slot: ' f'max_loras={capacity.max_loras}, used_loras={capacity.used_loras}, ' @@ -116,13 +122,16 @@ def main() -> None: session_heartbeat_interval=10, ) try: + # ModelManagement registers its capacity lazily on the first model + # request. ``MultiLoraTransformersModel.__init__`` calls /create, + # ensuring get_capacity_info() reports 0/max_loras instead of 0/0. + model = MultiLoraTransformersModel(model_id=args.served_model) _wait_for_free_lora( client, timeout=args.capacity_wait_seconds, poll_interval=args.capacity_poll_seconds, ) - model = MultiLoraTransformersModel(model_id=args.served_model) model.add_adapter_to_model( f'dsv4_diag_{args.mode}', LoraConfig( From 6c55e3789e1bd64e2c3ea5166c397d8ea11b5ca3 Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Wed, 2 Sep 2026 11:12:42 +0800 Subject: [PATCH 22/28] wip --- .../probe_dsv4_4layer_logits_cli.py | 228 ++++++++++ .../probe_dsv4_4layer_logits_ray.py | 251 ++++++++++ cookbook/client/server/transformer/read_me.md | 429 ++++++++++++++++-- .../run_dsv4_0731_npu_2node_2npu.sh | 8 +- .../server_config_dsv4_4layer_diag_ep.yaml | 6 +- .../server_config_dsv4_4layer_diag_no_ep.yaml | 6 +- 6 files changed, 883 insertions(+), 45 deletions(-) create mode 100644 cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_cli.py create mode 100644 cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_ray.py diff --git a/cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_cli.py b/cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_cli.py new file mode 100644 index 000000000..71c83876b --- /dev/null +++ b/cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_cli.py @@ -0,0 +1,228 @@ +#!/usr/bin/env python3 +"""Run a four-layer DeepSeek-V4 logits probe directly with torchrun. + +This is the local/CLI counterpart of ``probe_dsv4_4layer_logits.py``. It +constructs ``TransformersModel`` in every torchrun process and therefore does +not use GatewayServer, Ray Serve, sessions, tenants, or LoRA slots. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +from collections.abc import Mapping +from pathlib import Path +from typing import Any + + +DEFAULT_INPUT_IDS = [0, 128803, 2788, 6573, 70979, 36005, 320, 128804, 128821] + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--mode', required=True, choices=('no_ep', 'ep_loop', 'ep_gmm')) + parser.add_argument( + '--model-id', + default=os.environ.get('DSV4_MODEL_ID'), + help='Four-layer BF16 checkpoint. Defaults to DSV4_MODEL_ID.', + ) + parser.add_argument('--output-dir', type=Path, default=Path('/nas/disk6/ljl/dsv4_ep_diag/results')) + parser.add_argument('--input-ids', default=','.join(str(item) for item in DEFAULT_INPUT_IDS)) + parser.add_argument('--seed', type=int, default=42) + parser.add_argument('--mixed-precision', choices=('no', 'fp16', 'bf16'), default='bf16') + parser.add_argument( + '--ep-size', + type=int, + default=None, + help='EP size for ep_loop/ep_gmm. Defaults to the torchrun world size.', + ) + parser.add_argument( + '--disable-memory-efficient-init', + action='store_true', + help='Load the complete checkpoint in every process before FSDP wrapping.', + ) + return parser.parse_args() + + +def _extract_logits(result: Any): + import torch + + if hasattr(result, 'to_dict'): + result = result.to_dict() + elif hasattr(result, 'model_dump'): + result = result.model_dump() + if isinstance(result, list) and len(result) == 1 and isinstance(result[0], Mapping): + result = result[0] + if not isinstance(result, Mapping) or result.get('logits') is None: + raise RuntimeError(f'forward_only did not return logits; result type={type(result).__name__}') + + logits = result['logits'] + if not isinstance(logits, torch.Tensor): + logits = torch.as_tensor(logits) + logits = logits.detach().to(dtype=torch.float32) + original_shape = tuple(logits.shape) + while logits.ndim > 3: + logits = logits[0] + if logits.ndim == 3: + logits = logits[0, -1] + elif logits.ndim == 2: + logits = logits[-1] + elif logits.ndim != 1: + raise RuntimeError(f'Unsupported logits shape: {original_shape}') + if logits.numel() < 1000: + raise RuntimeError(f'Last-token logits look too small: shape={tuple(logits.shape)}') + return logits.contiguous() + + +def _tensor_sha256(tensor) -> str: + values = tensor.detach().cpu().contiguous().numpy().tobytes() + return hashlib.sha256(values).hexdigest() + + +def _check_rank_consistency(last_logits): + """Return the largest rank-to-rank difference without gathering Python objects.""" + import torch + import torch.distributed as dist + + if not dist.is_available() or not dist.is_initialized() or dist.get_world_size() == 1: + return 0.0 + + rank0_logits = last_logits.clone() + dist.broadcast(rank0_logits, src=0) + max_diff = (last_logits - rank0_logits).abs().max() + dist.all_reduce(max_diff, op=dist.ReduceOp.MAX) + return float(max_diff.item()) + + +def main() -> None: + args = parse_args() + if not args.model_id: + raise SystemExit('Set DSV4_MODEL_ID or pass --model-id.') + + input_ids = [int(item.strip()) for item in args.input_ids.split(',') if item.strip()] + if not input_ids: + raise SystemExit('--input-ids must contain at least one token ID') + + world_size = int(os.environ.get('WORLD_SIZE', '1')) + rank = int(os.environ.get('RANK', '0')) + enable_ep = args.mode != 'no_ep' + ep_size = args.ep_size if args.ep_size is not None else (world_size if enable_ep else 1) + if enable_ep and world_size <= 1: + raise SystemExit(f'{args.mode} requires at least two torchrun processes; WORLD_SIZE={world_size}.') + if enable_ep and (ep_size <= 1 or world_size % ep_size != 0): + raise SystemExit(f'Invalid EP topology: WORLD_SIZE={world_size}, ep_size={ep_size}.') + if not enable_ep and args.ep_size not in (None, 1): + raise SystemExit('no_ep only supports --ep-size 1.') + + # Kernel selection is read while EP patches are installed. Set it before + # importing Twinkle model modules so CLI and server modes choose the same + # implementation. + os.environ['TWINKLE_EP_FORCE_LOOP'] = '1' if args.mode == 'ep_loop' else '0' + os.environ.setdefault('TWINKLE_EP_DIAGNOSTICS', '1') + os.environ.setdefault('TWINKLE_TRUST_REMOTE_CODE', '1') + + import torch + import torch.distributed as dist + from transformers import AutoConfig + + import twinkle + from twinkle import DeviceMesh, Platform + from twinkle.model import TransformersModel + + device_mesh = DeviceMesh.from_sizes( + fsdp_size=world_size, + dp_size=1, + ep_size=ep_size, + device_type=Platform.get_platform().device_prefix(), + ) + twinkle.initialize( + mode='local', + global_device_mesh=device_mesh, + seed=args.seed, + lazy_collect=False, + ) + + config = AutoConfig.from_pretrained(args.model_id, trust_remote_code=True) + model = TransformersModel( + model_id=args.model_id, + config=config, + device_mesh=device_mesh, + strategy='native_fsdp', + mixed_precision=args.mixed_precision, + memory_efficient_init=not args.disable_memory_efficient_init, + fsdp_config={ + 'reshard_after_forward': True, + 'expert_parallel': { + 'enabled': enable_ep, + 'ep_size': ep_size, + 'router_dtype': 'fp32', + 'keep_router_logits': False, + }, + }, + ) + model.set_processor('InputProcessor', padding_side='left', padding_free=False) + + raw_input = { + 'input_ids': input_ids, + 'attention_mask': [1] * len(input_ids), + 'position_ids': list(range(len(input_ids))), + } + response = model.forward_only( + inputs=[raw_input], + return_logits=True, + ) + last_logits = _extract_logits(response) + max_rank_diff = _check_rank_consistency(last_logits) + last_logits_cpu = last_logits.cpu() + + if rank == 0: + finite = torch.isfinite(last_logits_cpu) + top_values, top_indices = torch.topk(last_logits_cpu, k=min(20, last_logits_cpu.numel())) + report = { + 'execution': 'cli', + 'mode': args.mode, + 'model_id': str(Path(args.model_id).expanduser()), + 'world_size': world_size, + 'ep_size': ep_size, + 'memory_efficient_init': not args.disable_memory_efficient_init, + 'input_ids': input_ids, + 'last_logits_shape': list(last_logits_cpu.shape), + 'dtype_saved': str(last_logits_cpu.dtype), + 'sha256': _tensor_sha256(last_logits_cpu), + 'max_rank_diff': max_rank_diff, + 'finite': bool(finite.all().item()), + 'nan_count': int(torch.isnan(last_logits_cpu).sum().item()), + 'inf_count': int(torch.isinf(last_logits_cpu).sum().item()), + 'sum': last_logits_cpu.sum().item(), + 'abs_sum': last_logits_cpu.abs().sum().item(), + 'min': last_logits_cpu.min().item(), + 'max': last_logits_cpu.max().item(), + 'top_token_ids': top_indices.tolist(), + 'top_logits': top_values.tolist(), + } + + args.output_dir.mkdir(parents=True, exist_ok=True) + tensor_path = args.output_dir / f'cli_{args.mode}_last_logits.pt' + json_path = args.output_dir / f'cli_{args.mode}_last_logits.json' + torch.save( + { + 'execution': 'cli', + 'mode': args.mode, + 'input_ids': input_ids, + 'last_logits': last_logits_cpu, + }, + tensor_path, + ) + json_path.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding='utf-8') + print(json.dumps(report, ensure_ascii=False, indent=2), flush=True) + print(f'Logits saved to: {tensor_path.resolve()}', flush=True) + print(f'Report saved to: {json_path.resolve()}', flush=True) + + if dist.is_available() and dist.is_initialized(): + dist.barrier() + + +if __name__ == '__main__': + main() diff --git a/cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_ray.py b/cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_ray.py new file mode 100644 index 000000000..c48269210 --- /dev/null +++ b/cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_ray.py @@ -0,0 +1,251 @@ +#!/usr/bin/env python3 +"""Run the four-layer DeepSeek-V4 logits probe on a multi-node Ray cluster. + +The script is a command-line Ray driver. It creates distributed Twinkle model +actors directly and deliberately does not deploy Ray Serve, GatewayServer, +ModelManagement, sessions, tenants, or LoRA adapters. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +from collections.abc import Mapping +from pathlib import Path +from typing import Any + + +DEFAULT_INPUT_IDS = [0, 128803, 2788, 6573, 70979, 36005, 320, 128804, 128821] + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--mode', required=True, choices=('no_ep', 'ep_loop', 'ep_gmm')) + parser.add_argument( + '--model-id', + default=os.environ.get('DSV4_MODEL_ID'), + help='Four-layer BF16 checkpoint. Defaults to DSV4_MODEL_ID.', + ) + parser.add_argument('--ray-address', default=os.environ.get('RAY_ADDRESS', 'auto')) + parser.add_argument('--world-size', type=int, default=4) + parser.add_argument('--nproc-per-node', type=int, default=2) + parser.add_argument('--output-dir', type=Path, default=Path('/nas/disk6/ljl/dsv4_ep_diag/results')) + parser.add_argument('--input-ids', default=','.join(str(item) for item in DEFAULT_INPUT_IDS)) + parser.add_argument('--seed', type=int, default=42) + parser.add_argument('--mixed-precision', choices=('no', 'fp16', 'bf16'), default='bf16') + parser.add_argument( + '--ep-size', + type=int, + default=None, + help='EP size for ep_loop/ep_gmm. Defaults to world-size.', + ) + parser.add_argument( + '--disable-memory-efficient-init', + action='store_true', + help='Load the complete checkpoint in every actor before FSDP wrapping.', + ) + return parser.parse_args() + + +def _extract_logits(result: Any): + import torch + + if callable(result): + result = result() + if hasattr(result, 'to_dict'): + result = result.to_dict() + elif hasattr(result, 'model_dump'): + result = result.model_dump() + if isinstance(result, list) and len(result) == 1 and isinstance(result[0], Mapping): + result = result[0] + if not isinstance(result, Mapping) or result.get('logits') is None: + raise RuntimeError(f'forward_only did not return logits; result type={type(result).__name__}') + + logits = result['logits'] + if not isinstance(logits, torch.Tensor): + logits = torch.as_tensor(logits) + logits = logits.detach().cpu().to(dtype=torch.float32) + original_shape = tuple(logits.shape) + while logits.ndim > 3: + logits = logits[0] + if logits.ndim == 3: + logits = logits[0, -1] + elif logits.ndim == 2: + logits = logits[-1] + elif logits.ndim != 1: + raise RuntimeError(f'Unsupported logits shape: {original_shape}') + if logits.numel() < 1000: + raise RuntimeError(f'Last-token logits look too small: shape={tuple(logits.shape)}') + return logits.contiguous() + + +def _tensor_sha256(tensor) -> str: + values = tensor.detach().cpu().contiguous().numpy().tobytes() + return hashlib.sha256(values).hexdigest() + + +def _validate_topology(args: argparse.Namespace, ray) -> None: + required_nodes = args.world_size // args.nproc_per_node + alive_nodes = [node for node in ray.nodes() if node.get('Alive', True)] + npu_nodes = [ + node for node in alive_nodes + if int(node.get('Resources', {}).get('NPU', 0)) >= args.nproc_per_node + ] + total_npus = int(ray.cluster_resources().get('NPU', 0)) + if total_npus < args.world_size or len(npu_nodes) < required_nodes: + raise RuntimeError( + 'Ray cluster does not have the requested NPU topology: ' + f'required world_size={args.world_size}, nodes={required_nodes}, ' + f'nproc_per_node={args.nproc_per_node}; found NPU={total_npus}, ' + f'eligible_nodes={len(npu_nodes)}. Start every Ray node with ' + f'--resources=\'{{"NPU": {args.nproc_per_node}}}\'.') + + +def main() -> None: + args = parse_args() + if not args.model_id: + raise SystemExit('Set DSV4_MODEL_ID or pass --model-id.') + if args.world_size <= 0 or args.nproc_per_node <= 0: + raise SystemExit('--world-size and --nproc-per-node must be positive.') + if args.world_size % args.nproc_per_node != 0: + raise SystemExit('--world-size must be divisible by --nproc-per-node.') + + input_ids = [int(item.strip()) for item in args.input_ids.split(',') if item.strip()] + if not input_ids: + raise SystemExit('--input-ids must contain at least one token ID') + + enable_ep = args.mode != 'no_ep' + ep_size = args.ep_size if args.ep_size is not None else (args.world_size if enable_ep else 1) + if enable_ep and (ep_size <= 1 or args.world_size % ep_size != 0): + raise SystemExit(f'Invalid EP topology: world_size={args.world_size}, ep_size={ep_size}.') + if not enable_ep and args.ep_size not in (None, 1): + raise SystemExit('no_ep only supports --ep-size 1.') + + # RayHelper copies the driver's environment into each model actor. Set EP + # selection before Twinkle creates the placement groups and runtime_envs. + os.environ['TWINKLE_EP_FORCE_LOOP'] = '1' if args.mode == 'ep_loop' else '0' + os.environ.setdefault('TWINKLE_EP_DIAGNOSTICS', '1') + os.environ.setdefault('TWINKLE_TRUST_REMOTE_CODE', '1') + os.environ.setdefault('RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES', '1') + + import ray + import torch + + ray.init(address=args.ray_address, ignore_reinit_error=True) + twinkle_initialized = False + try: + _validate_topology(args, ray) + + import twinkle + from twinkle import DeviceGroup, DeviceMesh + from twinkle.model import TransformersModel + + device_mesh = DeviceMesh.from_sizes( + fsdp_size=args.world_size, + dp_size=1, + ep_size=ep_size, + device_type='npu', + ) + groups = [ + DeviceGroup( + name='model', + ranks=list(range(args.world_size)), + device_type='NPU', + ) + ] + twinkle.initialize( + mode='ray', + nproc_per_node=args.nproc_per_node, + ncpu_proc_per_node=1, + seed=args.seed, + groups=groups, + global_device_mesh=device_mesh, + lazy_collect=False, + ) + twinkle_initialized = True + + model = TransformersModel( + model_id=args.model_id, + device_mesh=device_mesh, + remote_group='model', + instance_id=f'dsv4_4layer_ray_{args.mode}', + strategy='native_fsdp', + mixed_precision=args.mixed_precision, + memory_efficient_init=not args.disable_memory_efficient_init, + fsdp_config={ + 'reshard_after_forward': True, + 'expert_parallel': { + 'enabled': enable_ep, + 'ep_size': ep_size, + 'router_dtype': 'fp32', + 'keep_router_logits': False, + }, + }, + ) + model.set_processor('InputProcessor', padding_side='left', padding_free=False) + + raw_input = { + 'input_ids': input_ids, + 'attention_mask': [1] * len(input_ids), + 'position_ids': list(range(len(input_ids))), + } + response = model.forward_only( + inputs=[raw_input], + return_logits=True, + ) + last_logits = _extract_logits(response) + + finite = torch.isfinite(last_logits) + top_values, top_indices = torch.topk(last_logits, k=min(20, last_logits.numel())) + report = { + 'execution': 'ray_cli', + 'mode': args.mode, + 'ray_address': args.ray_address, + 'model_id': args.model_id, + 'world_size': args.world_size, + 'nproc_per_node': args.nproc_per_node, + 'ep_size': ep_size, + 'memory_efficient_init': not args.disable_memory_efficient_init, + 'input_ids': input_ids, + 'last_logits_shape': list(last_logits.shape), + 'dtype_saved': str(last_logits.dtype), + 'sha256': _tensor_sha256(last_logits), + 'finite': bool(finite.all().item()), + 'nan_count': int(torch.isnan(last_logits).sum().item()), + 'inf_count': int(torch.isinf(last_logits).sum().item()), + 'sum': last_logits.sum().item(), + 'abs_sum': last_logits.abs().sum().item(), + 'min': last_logits.min().item(), + 'max': last_logits.max().item(), + 'top_token_ids': top_indices.tolist(), + 'top_logits': top_values.tolist(), + } + + args.output_dir.mkdir(parents=True, exist_ok=True) + tensor_path = args.output_dir / f'ray_{args.mode}_last_logits.pt' + json_path = args.output_dir / f'ray_{args.mode}_last_logits.json' + torch.save( + { + 'execution': 'ray_cli', + 'mode': args.mode, + 'input_ids': input_ids, + 'last_logits': last_logits, + }, + tensor_path, + ) + json_path.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding='utf-8') + print(json.dumps(report, ensure_ascii=False, indent=2), flush=True) + print(f'Logits saved to: {tensor_path.resolve()}', flush=True) + print(f'Report saved to: {json_path.resolve()}', flush=True) + finally: + if twinkle_initialized: + from twinkle.infra._ray import RayHelper + + RayHelper.teardown() + ray.shutdown() + + +if __name__ == '__main__': + main() diff --git a/cookbook/client/server/transformer/read_me.md b/cookbook/client/server/transformer/read_me.md index 260279e8b..23de4287d 100644 --- a/cookbook/client/server/transformer/read_me.md +++ b/cookbook/client/server/transformer/read_me.md @@ -37,6 +37,8 @@ 诊断工具： - `cookbook/client/server/transformer/diagnostics/test_dsv4_npu_gmm_layout.py` +- `cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_ray.py` +- `cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_cli.py` - `cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits.py` - `cookbook/client/server/transformer/diagnostics/compare_dsv4_4layer_logits.py` @@ -71,7 +73,7 @@ hidden_size == 2 * moe_intermediate_size == 4096 ```bash PROJECT_DIR=/opt/twinkle -MODEL_DIR=/highcode/shared_data/DeepSeek-V4-Flash-0731-BF16-4layers +MODEL_DIR=/nas/disk6/ljl/DeepSeek-V4-Flash-0731-BF16-4layers cd "$PROJECT_DIR" git rev-parse HEAD @@ -96,16 +98,16 @@ Ray 的 Unix socket 路径必须短，但缓存又不应写入只有 10 GiB 的 ```bash mkdir -p /dev/shm/rh -mkdir -p /highcode/shared_data/dsv4_ep_diag/tmp -mkdir -p /highcode/shared_data/dsv4_ep_diag/logs -mkdir -p /highcode/shared_data/dsv4_ep_diag/results +mkdir -p /nas/disk6/ljl/dsv4_ep_diag/tmp +mkdir -p /nas/disk6/ljl/dsv4_ep_diag/logs +mkdir -p /nas/disk6/ljl/dsv4_ep_diag/results ``` 之后使用： ```bash export RAY_TMPDIR=/dev/shm/rh -export TMPDIR=/highcode/shared_data/dsv4_ep_diag/tmp +export TMPDIR=/nas/disk6/ljl/dsv4_ep_diag/tmp export RAY_ROTATION_MAX_BYTES=20971520 export RAY_ROTATION_BACKUP_COUNT=1 ``` @@ -131,7 +133,7 @@ ASCEND_RT_VISIBLE_DEVICES=0 \ PYTHONPATH=/opt/twinkle/src \ python3 cookbook/client/server/transformer/diagnostics/test_dsv4_npu_gmm_layout.py \ --device npu:0 \ - --output /highcode/shared_data/dsv4_ep_diag/results/npu_gmm_layout.json + --output /nas/disk6/ljl/dsv4_ep_diag/results/npu_gmm_layout.json ``` 成功时： @@ -144,6 +146,278 @@ old_bug_max_abs_diff 明显大于修复后的 max_abs_diff 如果这里失败，不要启动分布式服务，先处理 NPU GMM 与 `F.linear` 的数值差异。 +## 3A. 多机 Ray CLI 基准（推荐） + +这里的“CLI”表示在 Head 节点运行一个普通 Python driver，而不是必须使用 +torchrun。driver 直接通过 Ray 创建两节点 NPU 模型 Actor： + +```text +Python driver + -> Ray placement group + -> 两节点、每节点两个 TransformersModel Actor + -> native_fsdp + HCCL + EP + -> logits 返回 driver +``` + +该模式不启动 Ray Serve、GatewayServer 或 ModelManagement，不创建 HTTP session、 +租户或 LoRA adapter，因此不受槽位、心跳和接口限流影响；同时仍然覆盖 Ray 多机 +Actor、node-local 权重加载、`_broadcast_sharded_state_dict`、FSDP 和 EP。 + +先创建短临时目录和持久化结果目录： + +```bash +mkdir -p /dev/shm/rh +mkdir -p /nas/disk6/ljl/dsv4_ep_diag/tmp +mkdir -p /nas/disk6/ljl/dsv4_ep_diag/logs +mkdir -p /nas/disk6/ljl/dsv4_ep_diag/results +``` + +Head 节点： + +```bash +cd /opt/twinkle + +export HEAD_IP=172.61.10.254 +export ASCEND_RT_VISIBLE_DEVICES=0,1 +export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1 +export HCCL_SOCKET_IFNAME=eth0 +export GLOO_SOCKET_IFNAME=eth0 +export PYTHONPATH=/opt/twinkle/src +export TMPDIR=/nas/disk6/ljl/dsv4_ep_diag/tmp + +ray stop --force || true +ray start \ + --head \ + --node-ip-address="$HEAD_IP" \ + --port=6379 \ + --num-cpus="$(nproc)" \ + --resources='{"NPU": 2}' \ + --temp-dir=/dev/shm/rh \ + --disable-usage-stats \ + --include-dashboard=false +``` + +Worker 节点： + +```bash +cd /opt/twinkle + +export HEAD_IP=172.61.10.254 +export NODE_IP=172.61.12.251 +export ASCEND_RT_VISIBLE_DEVICES=0,1 +export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1 +export HCCL_SOCKET_IFNAME=eth0 +export GLOO_SOCKET_IFNAME=eth0 +export PYTHONPATH=/opt/twinkle/src +export TMPDIR=/nas/disk6/ljl/dsv4_ep_diag/tmp + +ray stop --force || true +ray start \ + --address="$HEAD_IP:6379" \ + --node-ip-address="$NODE_IP" \ + --num-cpus="$(nproc)" \ + --resources='{"NPU": 2}' \ + --temp-dir=/dev/shm/rh \ + --disable-usage-stats +``` + +Worker 加入后，只在 Head 节点运行 driver： + +```bash +cd /opt/twinkle + +export DSV4_MODEL_ID=/nas/disk6/ljl/DeepSeek-V4-Flash-0731-BF16-4layers +export ASCEND_RT_VISIBLE_DEVICES=0,1 +export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1 +export HCCL_SOCKET_IFNAME=eth0 +export GLOO_SOCKET_IFNAME=eth0 +export PYTHONPATH=/opt/twinkle/src +export TMPDIR=/nas/disk6/ljl/dsv4_ep_diag/tmp + +python3 cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_ray.py \ + --ray-address auto \ + --world-size 4 \ + --nproc-per-node 2 \ + --mode no_ep \ + --output-dir /nas/disk6/ljl/dsv4_ep_diag/results \ + 2>&1 | tee /nas/disk6/ljl/dsv4_ep_diag/logs/ray_no_ep.log +``` + +完成后，Ray 集群保持运行。分别重新执行 driver，只修改模式和日志名： + +```bash +python3 cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_ray.py \ + --ray-address auto --world-size 4 --nproc-per-node 2 \ + --mode ep_loop \ + --output-dir /nas/disk6/ljl/dsv4_ep_diag/results \ + 2>&1 | tee /nas/disk6/ljl/dsv4_ep_diag/logs/ray_ep_loop.log + +python3 cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_ray.py \ + --ray-address auto --world-size 4 --nproc-per-node 2 \ + --mode ep_gmm \ + --output-dir /nas/disk6/ljl/dsv4_ep_diag/results \ + 2>&1 | tee /nas/disk6/ljl/dsv4_ep_diag/logs/ray_ep_gmm.log +``` + +每个 driver 结束时只删除自己创建的模型 Actor 和 placement group，不执行 +`ray stop`。输出文件为： + +```text +ray_no_ep_last_logits.pt/json +ray_ep_loop_last_logits.pt/json +ray_ep_gmm_last_logits.pt/json +``` + +比较方式： + +```bash +python3 cookbook/client/server/transformer/diagnostics/compare_dsv4_4layer_logits.py \ + /nas/disk6/ljl/dsv4_ep_diag/results/ray_no_ep_last_logits.pt \ + /nas/disk6/ljl/dsv4_ep_diag/results/ray_ep_loop_last_logits.pt \ + --output /nas/disk6/ljl/dsv4_ep_diag/results/compare_ray_no_ep_vs_ep_loop.json + +python3 cookbook/client/server/transformer/diagnostics/compare_dsv4_4layer_logits.py \ + /nas/disk6/ljl/dsv4_ep_diag/results/ray_ep_loop_last_logits.pt \ + /nas/disk6/ljl/dsv4_ep_diag/results/ray_ep_gmm_last_logits.pt \ + --output /nas/disk6/ljl/dsv4_ep_diag/results/compare_ray_ep_loop_vs_ep_gmm.json +``` + +## 3B. 可选的本地/torchrun 基准 + +CLI 模式直接在 torchrun 进程中构造 `TransformersModel`，不启动 Ray、Gateway +或 ModelManagement，也不会创建 session、租户和 LoRA 槽位。该步骤验证： + +```text +Transformers checkpoint 转换 +native_fsdp 包装和权重加载 +EP 专家切分与通信 +loop/GMM 专家计算 +``` + +它不会验证 Ray Actor 和服务端 `_broadcast_sharded_state_dict` 路径，因此 CLI +通过后仍需执行后面的 C/S 对照。 + +### 单节点四张 NPU + +```bash +cd /opt/twinkle + +export DSV4_MODEL_ID=/nas/disk6/ljl/DeepSeek-V4-Flash-0731-BF16-4layers +export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3 +export HCCL_SOCKET_IFNAME=bond0 +export GLOO_SOCKET_IFNAME=bond0 +export PYTHONPATH=/opt/twinkle/src + +torchrun --standalone --nproc-per-node=4 \ + cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_cli.py \ + --mode no_ep \ + --output-dir /nas/disk6/ljl/dsv4_ep_diag/results + +torchrun --standalone --nproc-per-node=4 \ + cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_cli.py \ + --mode ep_loop \ + --output-dir /nas/disk6/ljl/dsv4_ep_diag/results + +torchrun --standalone --nproc-per-node=4 \ + cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_cli.py \ + --mode ep_gmm \ + --output-dir /nas/disk6/ljl/dsv4_ep_diag/results +``` + +每种模式必须启动一组新的 torchrun 进程，不能在同一 Python 进程内切换 +`TWINKLE_EP_FORCE_LOOP`。 + +### 两节点、每节点两张 NPU + +以下示例中 Head 为 `172.61.10.254`，Worker 为 `172.61.12.251`。先在 Head +执行 rank 0 命令，它会等待 Worker；然后在 Worker 执行 rank 1 命令。 + +Head： + +```bash +cd /opt/twinkle + +export DSV4_MODEL_ID=/nas/disk6/ljl/DeepSeek-V4-Flash-0731-BF16-4layers +export ASCEND_RT_VISIBLE_DEVICES=0,1 +export HCCL_SOCKET_IFNAME=eth0 +export GLOO_SOCKET_IFNAME=eth0 +export PYTHONPATH=/opt/twinkle/src + +torchrun \ + --nnodes=2 \ + --nproc-per-node=2 \ + --node-rank=0 \ + --master-addr=172.61.10.254 \ + --master-port=29610 \ + cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_cli.py \ + --mode no_ep \ + --output-dir /nas/disk6/ljl/dsv4_ep_diag/results +``` + +Worker： + +```bash +cd /opt/twinkle + +export DSV4_MODEL_ID=/nas/disk6/ljl/DeepSeek-V4-Flash-0731-BF16-4layers +export ASCEND_RT_VISIBLE_DEVICES=0,1 +export HCCL_SOCKET_IFNAME=eth0 +export GLOO_SOCKET_IFNAME=eth0 +export PYTHONPATH=/opt/twinkle/src + +torchrun \ + --nnodes=2 \ + --nproc-per-node=2 \ + --node-rank=1 \ + --master-addr=172.61.10.254 \ + --master-port=29610 \ + cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_cli.py \ + --mode no_ep \ + --output-dir /nas/disk6/ljl/dsv4_ep_diag/results +``` + +完成 `no_ep` 后，把两个节点命令中的 `--mode no_ep` 同时改成 +`--mode ep_loop` 再运行一次，然后改成 `--mode ep_gmm`。上一组 torchrun 已经 +完全退出后可以复用端口 `29610`；如果仍有残留进程，则先终止残留进程或换一个 +未占用端口。 + +只使用一张 NPU 时只能验证 `no_ep`： + +```bash +ASCEND_RT_VISIBLE_DEVICES=0 \ +PYTHONPATH=/opt/twinkle/src \ +python3 cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_cli.py \ + --mode no_ep \ + --model-id /nas/disk6/ljl/DeepSeek-V4-Flash-0731-BF16-4layers \ + --output-dir /nas/disk6/ljl/dsv4_ep_diag/results +``` + +CLI 输出： + +```text +cli_no_ep_last_logits.pt/json +cli_ep_loop_last_logits.pt/json +cli_ep_gmm_last_logits.pt/json +``` + +报告中的 `max_rank_diff` 应接近 0。使用现有比较脚本比较 CLI 三组结果： + +```bash +python3 cookbook/client/server/transformer/diagnostics/compare_dsv4_4layer_logits.py \ + /nas/disk6/ljl/dsv4_ep_diag/results/cli_no_ep_last_logits.pt \ + /nas/disk6/ljl/dsv4_ep_diag/results/cli_ep_loop_last_logits.pt \ + --output /nas/disk6/ljl/dsv4_ep_diag/results/compare_cli_no_ep_vs_ep_loop.json + +python3 cookbook/client/server/transformer/diagnostics/compare_dsv4_4layer_logits.py \ + /nas/disk6/ljl/dsv4_ep_diag/results/cli_ep_loop_last_logits.pt \ + /nas/disk6/ljl/dsv4_ep_diag/results/cli_ep_gmm_last_logits.pt \ + --output /nas/disk6/ljl/dsv4_ep_diag/results/compare_cli_ep_loop_vs_ep_gmm.json +``` + +还可以将 `cli__last_logits.pt` 与 C/S 产生的 +`_last_logits.pt` 比较。CLI 一致而 C/S 不一致时，重点排查 Ray/服务端权重 +分发；CLI 本身已经不一致时，重点排查 FSDP/EP/GMM。 + ## 4. 三组服务端对照模式 三组模式只改变 EP 和专家计算方式： @@ -181,8 +455,8 @@ rank 3 / ep_rank 3: experts 192..255 以下示例使用： ```text -Head: 172.61.10.111 -Worker: 172.61.12.165 +Head: 172.61.10.254 +Worker: 172.61.12.251 网卡: eth0 ``` @@ -193,15 +467,15 @@ Head 节点执行： ```bash cd /opt/twinkle -export DSV4_MODEL_ID=/highcode/shared_data/DeepSeek-V4-Flash-0731-BF16-4layers -export HEAD_IP=172.61.10.111 -export WORKER_IP=172.61.12.165 -export NODE_IP=172.61.10.111 +export DSV4_MODEL_ID=/nas/disk6/ljl/DeepSeek-V4-Flash-0731-BF16-4layers +export HEAD_IP=172.61.10.254 +export WORKER_IP=172.61.12.251 +export NODE_IP=172.61.10.254 export NETWORK_IFACE=eth0 export ASCEND_RT_VISIBLE_DEVICES=0,1 export RESET_RAY=1 export RAY_TMPDIR=/dev/shm/rh -export TMPDIR=/highcode/shared_data/dsv4_ep_diag/tmp +export TMPDIR=/nas/disk6/ljl/dsv4_ep_diag/tmp export RAY_ROTATION_MAX_BYTES=20971520 export RAY_ROTATION_BACKUP_COUNT=1 ``` @@ -211,15 +485,15 @@ Worker 节点执行： ```bash cd /opt/twinkle -export DSV4_MODEL_ID=/highcode/shared_data/DeepSeek-V4-Flash-0731-BF16-4layers -export HEAD_IP=172.61.10.111 -export WORKER_IP=172.61.12.165 -export NODE_IP=172.61.12.165 +export DSV4_MODEL_ID=/nas/disk6/ljl/DeepSeek-V4-Flash-0731-BF16-4layers +export HEAD_IP=172.61.10.254 +export WORKER_IP=172.61.12.251 +export NODE_IP=172.61.12.251 export NETWORK_IFACE=eth0 export ASCEND_RT_VISIBLE_DEVICES=0,1 export RESET_RAY=1 export RAY_TMPDIR=/dev/shm/rh -export TMPDIR=/highcode/shared_data/dsv4_ep_diag/tmp +export TMPDIR=/nas/disk6/ljl/dsv4_ep_diag/tmp export RAY_ROTATION_MAX_BYTES=20971520 export RAY_ROTATION_BACKUP_COUNT=1 ``` @@ -232,7 +506,7 @@ export RAY_ROTATION_BACKUP_COUNT=1 ```bash nohup bash cookbook/client/server/transformer/run_dsv4_4layer_ep_diagnostic.sh head no_ep \ - >/highcode/shared_data/dsv4_ep_diag/logs/no_ep_head.log 2>&1 /nas/disk6/ljl/dsv4_ep_diag/logs/no_ep_head.log 2>&1 /highcode/shared_data/dsv4_ep_diag/logs/no_ep_worker.log 2>&1 /nas/disk6/ljl/dsv4_ep_diag/logs/no_ep_worker.log 2>&1 /highcode/shared_data/dsv4_ep_diag/logs/ep_loop_head.log 2>&1 /nas/disk6/ljl/dsv4_ep_diag/logs/ep_loop_head.log 2>&1 /highcode/shared_data/dsv4_ep_diag/logs/ep_loop_worker.log 2>&1 /nas/disk6/ljl/dsv4_ep_diag/logs/ep_loop_worker.log 2>&1 /highcode/shared_data/dsv4_ep_diag/logs/ep_gmm_head.log 2>&1 /nas/disk6/ljl/dsv4_ep_diag/logs/ep_gmm_head.log 2>&1 /highcode/shared_data/dsv4_ep_diag/logs/ep_gmm_worker.log 2>&1 /nas/disk6/ljl/dsv4_ep_diag/logs/ep_gmm_worker.log 2>&1 /highcode/shared_data/dsv4_ep_diag/results/ep_diagnostics.txt +grep -h '\[EP_DIAG\]' /nas/disk6/ljl/dsv4_ep_diag/logs/ep_*_head.log \ + >/nas/disk6/ljl/dsv4_ep_diag/results/ep_diagnostics.txt ``` 重点检查以下日志。 @@ -492,6 +788,8 @@ python3 -m pytest -q tests/kernel/ops/test_moe.py python3 -m py_compile \ cookbook/client/server/transformer/diagnostics/test_dsv4_npu_gmm_layout.py \ + cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_ray.py \ + cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_cli.py \ cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits.py \ cookbook/client/server/transformer/diagnostics/compare_dsv4_4layer_logits.py @@ -539,8 +837,65 @@ ray stop --force 结果位于共享卷： ```text -/highcode/shared_data/dsv4_ep_diag/results -/highcode/shared_data/dsv4_ep_diag/logs +/nas/disk6/ljl/dsv4_ep_diag/results +/nas/disk6/ljl/dsv4_ep_diag/logs ``` 确认不再需要后可手动清理共享卷中的诊断日志和结果。`/dev/shm/rh` 仅用于当前 Pod 的 Ray 临时文件，Pod 删除后不会保留。 + +## 15. 路由 Management 类型 NameError + +如果应用注册路由时出现以下任意错误： + +```text +NameError: name 'GatewayServer' is not defined +NameError: name 'ModelManagement' is not defined +NameError: name 'SamplerManagement' is not defined +NameError: name 'ProcessorManagement' is not defined +``` + +说明运行版本中的 handler 只在 `TYPE_CHECKING` 分支导入了对应的 Management +类型，但 FastAPI 在运行时解析了路由参数的类型注解。这发生在服务配置加载阶段， +与模型权重、FSDP 和 EP 无关。 + +当前源码已在 Gateway、Model、Sampler 和 Processor 的 handler 中提供无循环导入 +的运行时类型别名，包括： + +```text +src/twinkle/server/gateway/openai_handlers.py +src/twinkle/server/gateway/tinker_handlers.py +src/twinkle/server/gateway/twinkle_handlers.py +src/twinkle/server/model/tinker_handlers.py +src/twinkle/server/model/twinkle_handlers.py +src/twinkle/server/sampler/tinker_handlers.py +src/twinkle/server/sampler/twinkle_handlers.py +src/twinkle/server/processor/twinkle_handlers.py +``` + +将修复后的源码同步到 head 和 worker 使用的相同 Twinkle 目录，然后重新启动服务。 +可先做静态检查： + +```bash +cd /opt/twinkle +python3 -m compileall -q \ + src/twinkle/server/gateway \ + src/twinkle/server/model \ + src/twinkle/server/sampler \ + src/twinkle/server/processor +``` + +如果错误仍然存在，先确认容器实际导入的源码位置： + +```bash +cd /opt/twinkle +PYTHONPATH=/opt/twinkle/src python3 - <<'PY' +import inspect +import twinkle.server.gateway.openai_handlers as module + +print(inspect.getfile(module)) +print('GatewayServer runtime alias:', module.GatewayServer) +PY +``` + +输出文件应来自 `/opt/twinkle/src/twinkle/server/gateway/`，且不应再次出现 +`NameError`。 diff --git a/cookbook/client/server/transformer/run_dsv4_0731_npu_2node_2npu.sh b/cookbook/client/server/transformer/run_dsv4_0731_npu_2node_2npu.sh index 814a9eb94..f7fadf5ca 100755 --- a/cookbook/client/server/transformer/run_dsv4_0731_npu_2node_2npu.sh +++ b/cookbook/client/server/transformer/run_dsv4_0731_npu_2node_2npu.sh @@ -1,10 +1,10 @@ #!/usr/bin/env bash set -euo pipefail -# Head node (172.61.10.111): +# Head node (172.61.10.254): # bash cookbook/client/server/transformer/run_dsv4_0731_npu_2node_2npu.sh head # -# Worker node (172.61.12.165): +# Worker node (172.61.12.251): # bash cookbook/client/server/transformer/run_dsv4_0731_npu_2node_2npu.sh worker # # Override NETWORK_IFACE when the Ray/HCCL interface is not eth0. Set @@ -16,8 +16,8 @@ if [[ "$ROLE" != "head" && "$ROLE" != "worker" ]]; then exit 2 fi -DEFAULT_HEAD_IP="172.61.10.111" -DEFAULT_WORKER_IP="172.61.12.165" +DEFAULT_HEAD_IP="172.61.10.254" +DEFAULT_WORKER_IP="172.61.12.251" HEAD_IP="${HEAD_IP:-$DEFAULT_HEAD_IP}" WORKER_IP="${WORKER_IP:-$DEFAULT_WORKER_IP}" diff --git a/cookbook/client/server/transformer/server_config_dsv4_4layer_diag_ep.yaml b/cookbook/client/server/transformer/server_config_dsv4_4layer_diag_ep.yaml index 13d9dabe6..cdbc82d76 100644 --- a/cookbook/client/server/transformer/server_config_dsv4_4layer_diag_ep.yaml +++ b/cookbook/client/server/transformer/server_config_dsv4_4layer_diag_ep.yaml @@ -68,8 +68,10 @@ applications: dp_size: 1 ep_size: 4 queue_config: - rps_limit: 1 - tps_limit: 10000 + # Functional diagnostics issue several serial setup requests in quick + # succession; do not let the production-style limiter mask EP results. + rps_limit: 100 + tps_limit: 100000 max_input_tokens: 512 queue_timeout: 7200 execution_timeout: 7200 diff --git a/cookbook/client/server/transformer/server_config_dsv4_4layer_diag_no_ep.yaml b/cookbook/client/server/transformer/server_config_dsv4_4layer_diag_no_ep.yaml index 3576f7794..8be2c5951 100644 --- a/cookbook/client/server/transformer/server_config_dsv4_4layer_diag_no_ep.yaml +++ b/cookbook/client/server/transformer/server_config_dsv4_4layer_diag_no_ep.yaml @@ -68,8 +68,10 @@ applications: dp_size: 1 ep_size: 1 queue_config: - rps_limit: 1 - tps_limit: 10000 + # Functional diagnostics issue several serial setup requests in quick + # succession; do not let the production-style limiter mask EP results. + rps_limit: 100 + tps_limit: 100000 max_input_tokens: 512 queue_timeout: 7200 execution_timeout: 7200 From 3e990f5cef257f2c80ecfa44d4e3e1bd2f5390d2 Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Wed, 2 Sep 2026 14:05:59 +0800 Subject: [PATCH 23/28] wip --- .gitignore | 1 + .../probe_dsv4_4layer_logits_cli.py | 38 ++++++++++++++--- .../probe_dsv4_4layer_logits_ray.py | 41 +++++++++++++++++-- cookbook/client/server/transformer/read_me.md | 30 +++++++++++--- 4 files changed, 95 insertions(+), 15 deletions(-) diff --git a/.gitignore b/.gitignore index 974b630a8..73882c1a0 100644 --- a/.gitignore +++ b/.gitignore @@ -150,6 +150,7 @@ megatron_output/ .qoder .kiro/ .claude/ +.codex_tmp/ # Pytorch *.pth diff --git a/cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_cli.py b/cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_cli.py index 71c83876b..d4b5bbd6b 100644 --- a/cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_cli.py +++ b/cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_cli.py @@ -2,8 +2,10 @@ """Run a four-layer DeepSeek-V4 logits probe directly with torchrun. This is the local/CLI counterpart of ``probe_dsv4_4layer_logits.py``. It -constructs ``TransformersModel`` in every torchrun process and therefore does -not use GatewayServer, Ray Serve, sessions, tenants, or LoRA slots. +constructs ``MultiLoraTransformersModel`` in every torchrun process and +therefore does not use GatewayServer, Ray Serve, sessions, tenants, or +server-side LoRA capacity. A disabled diagnostic adapter is installed to +match the production initialization path. """ from __future__ import annotations @@ -18,6 +20,7 @@ DEFAULT_INPUT_IDS = [0, 128803, 2788, 6573, 70979, 36005, 320, 128804, 128821] +TARGET_PARAMETERS = ['mlp.experts.gate_up_proj', 'mlp.experts.down_proj'] def parse_args() -> argparse.Namespace: @@ -125,11 +128,12 @@ def main() -> None: import torch import torch.distributed as dist + from peft import LoraConfig from transformers import AutoConfig import twinkle from twinkle import DeviceMesh, Platform - from twinkle.model import TransformersModel + from twinkle.model import MultiLoraTransformersModel device_mesh = DeviceMesh.from_sizes( fsdp_size=world_size, @@ -145,13 +149,18 @@ def main() -> None: ) config = AutoConfig.from_pretrained(args.model_id, trust_remote_code=True) - model = TransformersModel( + adapter_name = f'dsv4_diag_{args.mode}' + model = MultiLoraTransformersModel( model_id=args.model_id, config=config, device_mesh=device_mesh, strategy='native_fsdp', mixed_precision=args.mixed_precision, memory_efficient_init=not args.disable_memory_efficient_init, + max_loras=1, + max_r=8, + max_length=512, + target_modules='all-linear', fsdp_config={ 'reshard_after_forward': True, 'expert_parallel': { @@ -162,7 +171,24 @@ def main() -> None: }, }, ) - model.set_processor('InputProcessor', padding_side='left', padding_free=False) + model.add_adapter_to_model( + adapter_name, + LoraConfig( + r=8, + lora_alpha=32, + lora_dropout=0.0, + target_modules=None, + target_parameters=TARGET_PARAMETERS, + bias='none', + ), + gradient_accumulation_steps=1, + ) + model.set_processor( + 'InputProcessor', + adapter_name=adapter_name, + padding_side='left', + padding_free=False, + ) raw_input = { 'input_ids': input_ids, @@ -171,6 +197,8 @@ def main() -> None: } response = model.forward_only( inputs=[raw_input], + adapter_name=adapter_name, + disable_lora=True, return_logits=True, ) last_logits = _extract_logits(response) diff --git a/cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_ray.py b/cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_ray.py index c48269210..bed3e0214 100644 --- a/cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_ray.py +++ b/cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_ray.py @@ -3,7 +3,9 @@ The script is a command-line Ray driver. It creates distributed Twinkle model actors directly and deliberately does not deploy Ray Serve, GatewayServer, -ModelManagement, sessions, tenants, or LoRA adapters. +ModelManagement, sessions, tenants, or server-side LoRA capacity. A disabled +diagnostic adapter is installed inside the model to match the production +MultiLoraTransformersModel initialization path. """ from __future__ import annotations @@ -18,6 +20,7 @@ DEFAULT_INPUT_IDS = [0, 128803, 2788, 6573, 70979, 36005, 320, 128804, 128821] +TARGET_PARAMETERS = ['mlp.experts.gate_up_proj', 'mlp.experts.down_proj'] def parse_args() -> argparse.Namespace: @@ -138,9 +141,11 @@ def main() -> None: try: _validate_topology(args, ray) + from peft import LoraConfig + import twinkle from twinkle import DeviceGroup, DeviceMesh - from twinkle.model import TransformersModel + from twinkle.model import MultiLoraTransformersModel device_mesh = DeviceMesh.from_sizes( fsdp_size=args.world_size, @@ -166,7 +171,12 @@ def main() -> None: ) twinkle_initialized = True - model = TransformersModel( + # Match the model initialization used by ModelManagement. The + # Multi-LoRA constructor installs its preallocated slots and aligns + # their dtype before native FSDP wrapping. There is no server-side + # session or capacity accounting in this direct-Ray driver. + adapter_name = f'dsv4_diag_{args.mode}' + model = MultiLoraTransformersModel( model_id=args.model_id, device_mesh=device_mesh, remote_group='model', @@ -174,6 +184,10 @@ def main() -> None: strategy='native_fsdp', mixed_precision=args.mixed_precision, memory_efficient_init=not args.disable_memory_efficient_init, + max_loras=1, + max_r=8, + max_length=512, + target_modules='all-linear', fsdp_config={ 'reshard_after_forward': True, 'expert_parallel': { @@ -184,7 +198,24 @@ def main() -> None: }, }, ) - model.set_processor('InputProcessor', padding_side='left', padding_free=False) + model.add_adapter_to_model( + adapter_name, + LoraConfig( + r=8, + lora_alpha=32, + lora_dropout=0.0, + target_modules=None, + target_parameters=TARGET_PARAMETERS, + bias='none', + ), + gradient_accumulation_steps=1, + ) + model.set_processor( + 'InputProcessor', + adapter_name=adapter_name, + padding_side='left', + padding_free=False, + ) raw_input = { 'input_ids': input_ids, @@ -193,6 +224,8 @@ def main() -> None: } response = model.forward_only( inputs=[raw_input], + adapter_name=adapter_name, + disable_lora=True, return_logits=True, ) last_logits = _extract_logits(response) diff --git a/cookbook/client/server/transformer/read_me.md b/cookbook/client/server/transformer/read_me.md index 23de4287d..a0c369a4f 100644 --- a/cookbook/client/server/transformer/read_me.md +++ b/cookbook/client/server/transformer/read_me.md @@ -154,14 +154,17 @@ torchrun。driver 直接通过 Ray 创建两节点 NPU 模型 Actor： ```text Python driver -> Ray placement group - -> 两节点、每节点两个 TransformersModel Actor + -> 两节点、每节点两个 MultiLoraTransformersModel Actor -> native_fsdp + HCCL + EP -> logits 返回 driver ``` -该模式不启动 Ray Serve、GatewayServer 或 ModelManagement，不创建 HTTP session、 -租户或 LoRA adapter，因此不受槽位、心跳和接口限流影响；同时仍然覆盖 Ray 多机 -Actor、node-local 权重加载、`_broadcast_sharded_state_dict`、FSDP 和 EP。 +该模式不启动 Ray Serve、GatewayServer 或 ModelManagement，不创建 HTTP session +或租户，因此不受服务端槽位、心跳和接口限流影响。模型内部会创建一个禁用 delta +的诊断 adapter，以复用生产 `MultiLoraTransformersModel` 在 FSDP 包装前安装槽位和 +对齐 dtype 的相同初始化路径；该 adapter 不登记到 ServerState 容量中。该模式仍然 +覆盖 Ray 多机 Actor、node-local 权重加载、`_broadcast_sharded_state_dict`、FSDP +和 EP。 先创建短临时目录和持久化结果目录： @@ -268,6 +271,19 @@ ray_ep_loop_last_logits.pt/json ray_ep_gmm_last_logits.pt/json ``` +如果首次 forward 报错： + +```text +AssertionError: FSDP expects uniform original parameter dtype but got +{torch.bfloat16, torch.float32} +``` + +说明运行的是旧版 Ray probe：它实例化了纯 `TransformersModel`，没有进入生产 +Multi-LoRA 的槽位安装和 `_ensure_lora_dtype()` 路径。使用更新后的 +`probe_dsv4_4layer_logits_ray.py`；它直接构造 `MultiLoraTransformersModel`，先安装 +禁用的诊断 adapter，再进入 `_lazy_wrap_model()`。无需修改公共 +`TransformersModel._lazy_wrap_model()`，Ray 集群也无需重启。 + 比较方式： ```bash @@ -284,8 +300,10 @@ python3 cookbook/client/server/transformer/diagnostics/compare_dsv4_4layer_logit ## 3B. 可选的本地/torchrun 基准 -CLI 模式直接在 torchrun 进程中构造 `TransformersModel`，不启动 Ray、Gateway -或 ModelManagement，也不会创建 session、租户和 LoRA 槽位。该步骤验证： +CLI 模式直接在 torchrun 进程中构造 `MultiLoraTransformersModel`，不启动 Ray、 +Gateway 或 ModelManagement，也不会创建 session、租户或登记服务端 LoRA 容量。 +它会在每个进程内安装一个禁用 delta 的诊断 adapter，使 dtype 对齐和 FSDP 包装 +顺序与生产服务一致。该步骤验证： ```text Transformers checkpoint 转换 From aa966a0eb1068b51d41b41c2b74980d8c6e246cf Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Wed, 2 Sep 2026 16:36:56 +0800 Subject: [PATCH 24/28] wip --- .../diagnostics/compare_dsv4_4layer_logits.py | 28 +++++++++------- tests/kernel/ops/test_ep_dispatch.py | 33 +++++++++++++++++++ 2 files changed, 50 insertions(+), 11 deletions(-) create mode 100644 tests/kernel/ops/test_ep_dispatch.py diff --git a/cookbook/client/server/transformer/diagnostics/compare_dsv4_4layer_logits.py b/cookbook/client/server/transformer/diagnostics/compare_dsv4_4layer_logits.py index 58facf908..9fce9ef63 100755 --- a/cookbook/client/server/transformer/diagnostics/compare_dsv4_4layer_logits.py +++ b/cookbook/client/server/transformer/diagnostics/compare_dsv4_4layer_logits.py @@ -22,22 +22,29 @@ def parse_args() -> argparse.Namespace: def load_logits(path: Path) -> tuple[str, list[int], torch.Tensor]: payload = torch.load(path, map_location='cpu', weights_only=True) - return str(payload.get('mode', path.stem)), list(payload['input_ids']), payload['last_logits'].float() + return ( + str(payload.get('mode', path.stem)), + list(payload['input_ids']), + payload['last_logits'].float(), + ) def main() -> None: args = parse_args() reference_mode, reference_ids, reference = load_logits(args.reference) candidate_mode, candidate_ids, candidate = load_logits(args.candidate) + if reference_ids != candidate_ids: raise RuntimeError(f'Input IDs differ: {reference_ids} != {candidate_ids}') if reference.shape != candidate.shape: raise RuntimeError(f'Logit shapes differ: {tuple(reference.shape)} != {tuple(candidate.shape)}') - difference = (candidate - reference).abs() - close = torch.isclose(candidate, reference, rtol=args.rtol, atol=args.atol) - reference_top = torch.topk(reference, k=min(20, reference.numel())).indices.tolist() - candidate_top = torch.topk(candidate, k=min(20, candidate.numel())).indices.tolist() + difference = (reference - candidate).abs() + close = torch.isclose(reference, candidate, rtol=args.rtol, atol=args.atol) + top_k = min(20, reference.numel()) + reference_top = torch.topk(reference, k=top_k).indices.tolist() + candidate_top = torch.topk(candidate, k=top_k).indices.tolist() + report = { 'reference': str(args.reference), 'reference_mode': reference_mode, @@ -48,16 +55,15 @@ def main() -> None: 'rtol': args.rtol, 'atol': args.atol, 'allclose': bool(close.all().item()), - 'close_fraction': close.float().mean().item(), - 'max_abs_diff': difference.max().item(), - 'mean_abs_diff': difference.mean().item(), + 'close_fraction': float(close.float().mean().item()), + 'max_abs_diff': float(difference.max().item()), + 'mean_abs_diff': float(difference.mean().item()), 'reference_top20': reference_top, 'candidate_top20': candidate_top, - 'top20_overlap': len(set(reference_top) & set(candidate_top)), + 'top20_overlap': len(set(reference_top) & set(candidate_top)) / top_k, } - output = args.output or args.candidate.with_name( - f'compare_{reference_mode}_vs_{candidate_mode}.json') + output = args.output or args.reference.with_name(f'compare_{reference_mode}_vs_{candidate_mode}.json') output.parent.mkdir(parents=True, exist_ok=True) output.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding='utf-8') print(json.dumps(report, ensure_ascii=False, indent=2)) diff --git a/tests/kernel/ops/test_ep_dispatch.py b/tests/kernel/ops/test_ep_dispatch.py new file mode 100644 index 000000000..3fde6031f --- /dev/null +++ b/tests/kernel/ops/test_ep_dispatch.py @@ -0,0 +1,33 @@ +import torch + +from twinkle.kernel.ops import ep as ep_ops +from twinkle.kernel.ops.ep import loop as loop_ops + + +class _Backend(ep_ops.EpExpertsGmm): + + def __init__(self, name: str, value: float, *, fallback: bool = False): + self.name = name + self.value = value + self.fallback = fallback + + def ineligible_reason(self, experts_mod): + return None + + def forward(self, experts_mod, permuted_tokens, num_global_sum_tokens_per_local_expert, experts_per_rank): + return torch.full_like(permuted_tokens, self.value) + + +def test_ep_forward_uses_accelerated_backend_by_default(monkeypatch): + monkeypatch.delenv('TWINKLE_EP_FORCE_LOOP', raising=False) + monkeypatch.setattr(ep_ops, '_IMPLS', [_Backend('accelerated', 1.0), _Backend('loop', 2.0, fallback=True)]) + result = ep_ops.ep_forward(None, torch.zeros(2, 3), torch.tensor([2]), 1) + assert torch.equal(result, torch.ones(2, 3)) + + +def test_ep_forward_can_force_reference_loop(monkeypatch): + monkeypatch.setenv('TWINKLE_EP_FORCE_LOOP', '1') + monkeypatch.setattr(ep_ops, '_IMPLS', [_Backend('accelerated', 1.0), _Backend('loop', 2.0, fallback=True)]) + monkeypatch.setattr(loop_ops, 'LoopEpExpertsGmm', lambda: _Backend('loop', 2.0, fallback=True)) + result = ep_ops.ep_forward(None, torch.zeros(2, 3), torch.tensor([2]), 1) + assert torch.equal(result, torch.full((2, 3), 2.0)) From 40970bee37b58c7f271692dd6be3501e6b81ef98 Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Wed, 2 Sep 2026 17:23:07 +0800 Subject: [PATCH 25/28] wip --- src/twinkle/model/transformers/moe/ep_utils.py | 2 +- .../model/transformers/moe/expert_parallel.py | 14 +++++++++++--- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/src/twinkle/model/transformers/moe/ep_utils.py b/src/twinkle/model/transformers/moe/ep_utils.py index bb2bed884..d070859ba 100644 --- a/src/twinkle/model/transformers/moe/ep_utils.py +++ b/src/twinkle/model/transformers/moe/ep_utils.py @@ -190,7 +190,7 @@ def preprocess( expert_mask: torch.Tensor, num_experts: int, ep_group: dist.ProcessGroup, -) -> torch.Tensor: +) -> tuple[list[int], list[int], torch.Tensor, torch.Tensor]: ep_size = ep_group.size() num_local_experts = num_experts // ep_size rank = dist.get_rank(ep_group) diff --git a/src/twinkle/model/transformers/moe/expert_parallel.py b/src/twinkle/model/transformers/moe/expert_parallel.py index 0bc3a391d..867e10e89 100644 --- a/src/twinkle/model/transformers/moe/expert_parallel.py +++ b/src/twinkle/model/transformers/moe/expert_parallel.py @@ -21,6 +21,12 @@ def _ep_diagnostics_enabled() -> bool: return os.environ.get('TWINKLE_EP_DIAGNOSTICS', '').strip().lower() in {'1', 'true', 'yes', 'on'} +def _split_sizes_to_list(split_sizes: list[int] | torch.Tensor) -> list[int]: + if isinstance(split_sizes, torch.Tensor): + return [int(value) for value in split_sizes.detach().cpu().tolist()] + return [int(value) for value in split_sizes] + + @dataclass class ExpertParallelConfig: enabled: bool = True @@ -272,8 +278,10 @@ def forward(hidden_states: torch.Tensor, *args, **kwargs): log_diagnostics = _ep_diagnostics_enabled() and not getattr(block, '_ep_diagnostics_logged', False) if log_diagnostics: + input_splits_list = _split_sizes_to_list(input_splits) + output_splits_list = _split_sizes_to_list(output_splits) expected_assignments = hidden_states_2d.shape[0] * top_k - actual_assignments = int(input_splits.sum().item()) + actual_assignments = sum(input_splits_list) if actual_assignments != expected_assignments: raise RuntimeError( f'EP routing assignment mismatch for {block_name}: ' @@ -290,8 +298,8 @@ def forward(hidden_states: torch.Tensor, *args, **kwargs): block._ep_local_end, hidden_states_2d.shape[0], top_k, - input_splits.tolist(), - output_splits.tolist(), + input_splits_list, + output_splits_list, selected_preview, weights_preview, ) From 91f1a30c094758ab4d93834d403c307d89e452c5 Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Thu, 3 Sep 2026 08:44:27 +0800 Subject: [PATCH 26/28] wip --- .../compare_dsv4_weight_fingerprints.py | 210 ++++++++++++++++++ .../probe_dsv4_4layer_logits_cli.py | 144 ++++++++++++ cookbook/client/server/transformer/read_me.md | 66 +++++- 3 files changed, 419 insertions(+), 1 deletion(-) create mode 100755 cookbook/client/server/transformer/diagnostics/compare_dsv4_weight_fingerprints.py diff --git a/cookbook/client/server/transformer/diagnostics/compare_dsv4_weight_fingerprints.py b/cookbook/client/server/transformer/diagnostics/compare_dsv4_weight_fingerprints.py new file mode 100755 index 000000000..093f254eb --- /dev/null +++ b/cookbook/client/server/transformer/diagnostics/compare_dsv4_weight_fingerprints.py @@ -0,0 +1,210 @@ +#!/usr/bin/env python3 +"""Compare rank-local parameter fingerprints from two DSV4 CLI probe runs.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + + +METADATA_FIELDS = ( + 'global_shape', + 'global_stride', + 'local_shape', + 'local_stride', + 'dtype', + 'local_numel', + 'local_is_contiguous', + 'sample_count', +) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('reference_dir', type=Path) + parser.add_argument('candidate_dir', type=Path) + parser.add_argument('--mode', default='ep_loop', choices=('no_ep', 'ep_loop', 'ep_gmm')) + parser.add_argument('--max-mismatches', type=int, default=50) + parser.add_argument('--output', type=Path) + return parser.parse_args() + + +def load_rank_reports(directory: Path, mode: str) -> dict[int, dict[str, Any]]: + pattern = f'cli_{mode}_weight_fingerprints_rank*.json' + reports = {} + for path in sorted(directory.glob(pattern)): + payload = json.loads(path.read_text(encoding='utf-8')) + rank = int(payload['rank']) + if rank in reports: + raise RuntimeError(f'Duplicate rank {rank} in {directory}') + payload['_path'] = str(path) + reports[rank] = payload + if not reports: + raise FileNotFoundError(f'No files matching {directory / pattern}') + return reports + + +def parameter_map(report: dict[str, Any]) -> dict[str, dict[str, Any]]: + return {record['name']: record for record in report['parameters']} + + +def unravel_index(index: int, shape: list[int]) -> list[int]: + coordinate = [] + remaining = index + for size in reversed(shape): + coordinate.append(remaining % size) + remaining //= size + return list(reversed(coordinate)) + + +def sample_positions(numel: int, count: int) -> list[int]: + if numel <= 0 or count <= 0: + return [] + count = min(numel, count) + if count == 1: + return [0] + return [index * (numel - 1) // (count - 1) for index in range(count)] + + +def describe_value_mismatch(reference: dict[str, Any], candidate: dict[str, Any]) -> dict[str, Any]: + reference_values = reference['sample_values'] + candidate_values = candidate['sample_values'] + if len(reference_values) != len(candidate_values): + return { + 'reference_sample_count': len(reference_values), + 'candidate_sample_count': len(candidate_values), + } + + differences = [abs(float(left) - float(right)) for left, right in zip(reference_values, candidate_values)] + differing = [index for index, difference in enumerate(differences) if difference != 0.0] + result = { + 'different_sample_count': len(differing), + 'sample_count': len(differences), + 'max_abs_sample_diff': max(differences, default=0.0), + 'mean_abs_sample_diff': sum(differences) / len(differences) if differences else 0.0, + } + if differing: + sample_index = differing[0] + flat_positions = sample_positions(reference['local_numel'], reference['sample_count']) + flat_index = flat_positions[sample_index] + result['first_difference'] = { + 'sample_index': sample_index, + 'local_flat_index': flat_index, + 'local_coordinate': unravel_index(flat_index, reference['local_shape']), + 'reference': reference_values[sample_index], + 'candidate': candidate_values[sample_index], + } + return result + + +def main() -> None: + args = parse_args() + if args.max_mismatches <= 0: + raise SystemExit('--max-mismatches must be positive') + + reference_reports = load_rank_reports(args.reference_dir, args.mode) + candidate_reports = load_rank_reports(args.candidate_dir, args.mode) + reference_ranks = set(reference_reports) + candidate_ranks = set(candidate_reports) + if reference_ranks != candidate_ranks: + raise RuntimeError( + f'Rank sets differ: reference={sorted(reference_ranks)}, candidate={sorted(candidate_ranks)}') + + metadata_mismatches = [] + value_mismatches = [] + missing_parameters = [] + unexpected_parameters = [] + compared_parameters = 0 + exact_parameters = 0 + + for rank in sorted(reference_ranks): + reference_report = reference_reports[rank] + candidate_report = candidate_reports[rank] + if reference_report['sample_algorithm'] != candidate_report['sample_algorithm']: + raise RuntimeError(f'Rank {rank}: sample algorithms differ') + + reference_parameters = parameter_map(reference_report) + candidate_parameters = parameter_map(candidate_report) + reference_names = set(reference_parameters) + candidate_names = set(candidate_parameters) + for name in sorted(reference_names - candidate_names): + missing_parameters.append({'rank': rank, 'name': name}) + for name in sorted(candidate_names - reference_names): + unexpected_parameters.append({'rank': rank, 'name': name}) + + for name in sorted(reference_names & candidate_names): + compared_parameters += 1 + reference = reference_parameters[name] + candidate = candidate_parameters[name] + changed_metadata = { + field: {'reference': reference[field], 'candidate': candidate[field]} + for field in METADATA_FIELDS + if reference[field] != candidate[field] + } + if changed_metadata: + metadata_mismatches.append({ + 'rank': rank, + 'name': name, + 'fields': changed_metadata, + }) + continue + if reference['sample_sha256'] == candidate['sample_sha256']: + exact_parameters += 1 + continue + mismatch = { + 'rank': rank, + 'name': name, + 'reference_sha256': reference['sample_sha256'], + 'candidate_sha256': candidate['sample_sha256'], + } + mismatch.update(describe_value_mismatch(reference, candidate)) + value_mismatches.append(mismatch) + + total_mismatches = ( + len(metadata_mismatches) + + len(value_mismatches) + + len(missing_parameters) + + len(unexpected_parameters) + ) + report = { + 'mode': args.mode, + 'reference_dir': str(args.reference_dir), + 'candidate_dir': str(args.candidate_dir), + 'reference_memory_efficient_init': reference_reports[min(reference_ranks)]['memory_efficient_init'], + 'candidate_memory_efficient_init': candidate_reports[min(candidate_ranks)]['memory_efficient_init'], + 'ranks': sorted(reference_ranks), + 'compared_parameters': compared_parameters, + 'exact_sampled_parameters': exact_parameters, + 'exact_sampled_fraction': exact_parameters / compared_parameters if compared_parameters else 0.0, + 'metadata_mismatch_count': len(metadata_mismatches), + 'value_mismatch_count': len(value_mismatches), + 'missing_parameter_count': len(missing_parameters), + 'unexpected_parameter_count': len(unexpected_parameters), + 'passed': total_mismatches == 0, + 'metadata_mismatches': metadata_mismatches[:args.max_mismatches], + 'value_mismatches': value_mismatches[:args.max_mismatches], + 'missing_parameters': missing_parameters[:args.max_mismatches], + 'unexpected_parameters': unexpected_parameters[:args.max_mismatches], + 'mismatch_lists_truncated': any( + len(items) > args.max_mismatches + for items in ( + metadata_mismatches, + value_mismatches, + missing_parameters, + unexpected_parameters, + ) + ), + } + output = args.output or args.candidate_dir / f'compare_{args.mode}_weight_fingerprints.json' + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding='utf-8') + print(json.dumps(report, ensure_ascii=False, indent=2)) + print(f'Report saved to: {output.resolve()}') + if not report['passed']: + raise SystemExit(1) + + +if __name__ == '__main__': + main() diff --git a/cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_cli.py b/cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_cli.py index d4b5bbd6b..44684dac5 100644 --- a/cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_cli.py +++ b/cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_cli.py @@ -46,6 +46,17 @@ def parse_args() -> argparse.Namespace: action='store_true', help='Load the complete checkpoint in every process before FSDP wrapping.', ) + parser.add_argument( + '--dump-weight-fingerprints', + action='store_true', + help='Save sampled local base-parameter fingerprints for every rank after FSDP/EP wrapping.', + ) + parser.add_argument( + '--weight-fingerprint-samples', + type=int, + default=4096, + help='Maximum number of evenly spaced values sampled from each local parameter.', + ) return parser.parse_args() @@ -99,6 +110,126 @@ def _check_rank_consistency(last_logits): return float(max_diff.item()) +def _sample_positions(numel: int, count: int) -> list[int]: + """Return deterministic, evenly spaced flat indices without float rounding.""" + if numel <= 0 or count <= 0: + return [] + count = min(numel, count) + if count == 1: + return [0] + return [index * (numel - 1) // (count - 1) for index in range(count)] + + +def _parameter_local_tensor(parameter): + """Return the rank-local tensor for both regular Parameters and FSDP2 DTensors.""" + from torch.distributed.tensor import DTensor + + value = parameter.detach() + if isinstance(value, DTensor): + value = value.to_local() + return value + + +def _sample_tensor_at_flat_positions(tensor, positions: list[int]): + """Sample logical flat positions without flattening/copying the full tensor.""" + import torch + + if not positions: + return torch.empty(0, dtype=tensor.dtype, device=tensor.device) + if tensor.ndim == 0: + return tensor.reshape(1) + + remaining = positions + coordinate_columns = [] + for size in reversed(tensor.shape): + coordinate_columns.append([position % size for position in remaining]) + remaining = [position // size for position in remaining] + coordinates = tuple( + torch.tensor(column, dtype=torch.long, device=tensor.device) + for column in reversed(coordinate_columns) + ) + return tensor[coordinates] + + +def _sample_parameter(parameter, sample_limit: int) -> dict[str, Any]: + import torch + + local = _parameter_local_tensor(parameter) + positions = _sample_positions(local.numel(), sample_limit) + if positions: + sampled = _sample_tensor_at_flat_positions(local, positions).detach().cpu().contiguous() + # Hash the original dtype bytes. Converting BF16 directly to NumPy is + # not supported by every NumPy version, whereas a byte view is. + sample_sha256 = hashlib.sha256(sampled.view(torch.uint8).numpy().tobytes()).hexdigest() + sampled_float = sampled.float() + sample_values = sampled_float.tolist() + sample_finite = bool(torch.isfinite(sampled_float).all().item()) + sample_min = float(sampled_float.min().item()) + sample_max = float(sampled_float.max().item()) + sample_sum = float(sampled_float.sum().item()) + sample_abs_sum = float(sampled_float.abs().sum().item()) + else: + sample_sha256 = hashlib.sha256(b'').hexdigest() + sample_values = [] + sample_finite = True + sample_min = None + sample_max = None + sample_sum = 0.0 + sample_abs_sum = 0.0 + + return { + 'global_shape': list(parameter.shape), + 'global_stride': list(parameter.stride()), + 'local_shape': list(local.shape), + 'local_stride': list(local.stride()), + 'dtype': str(local.dtype), + 'local_numel': local.numel(), + 'local_is_contiguous': local.is_contiguous(), + 'sample_count': len(positions), + 'sample_sha256': sample_sha256, + 'sample_finite': sample_finite, + 'sample_sum': sample_sum, + 'sample_abs_sum': sample_abs_sum, + 'sample_min': sample_min, + 'sample_max': sample_max, + 'sample_values': sample_values, + } + + +def _write_weight_fingerprints(model, args: argparse.Namespace, *, rank: int, world_size: int) -> Path: + """Write rank-local post-wrap base-weight fingerprints to a JSON file.""" + import torch + + inner_model = getattr(model, 'model', model) + parameters = [] + with torch.no_grad(): + for name, parameter in inner_model.named_parameters(): + # The probe disables LoRA during forward. Adapter slots are seeded + # separately and are not part of the base-weight loading question. + if 'lora_' in name: + continue + record = {'name': name} + record.update(_sample_parameter(parameter, args.weight_fingerprint_samples)) + parameters.append(record) + + payload = { + 'execution': 'cli', + 'mode': args.mode, + 'model_id': str(Path(args.model_id).expanduser()), + 'rank': rank, + 'world_size': world_size, + 'memory_efficient_init': not args.disable_memory_efficient_init, + 'sample_algorithm': 'flat_evenly_spaced_integer_v1', + 'sample_limit_per_parameter': args.weight_fingerprint_samples, + 'parameter_count': len(parameters), + 'parameters': parameters, + } + args.output_dir.mkdir(parents=True, exist_ok=True) + output = args.output_dir / f'cli_{args.mode}_weight_fingerprints_rank{rank}.json' + output.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding='utf-8') + return output + + def main() -> None: args = parse_args() if not args.model_id: @@ -107,6 +238,8 @@ def main() -> None: input_ids = [int(item.strip()) for item in args.input_ids.split(',') if item.strip()] if not input_ids: raise SystemExit('--input-ids must contain at least one token ID') + if args.weight_fingerprint_samples <= 0: + raise SystemExit('--weight-fingerprint-samples must be positive') world_size = int(os.environ.get('WORLD_SIZE', '1')) rank = int(os.environ.get('RANK', '0')) @@ -202,6 +335,14 @@ def main() -> None: return_logits=True, ) last_logits = _extract_logits(response) + fingerprint_path = None + if args.dump_weight_fingerprints: + fingerprint_path = _write_weight_fingerprints( + model, + args, + rank=rank, + world_size=world_size, + ) max_rank_diff = _check_rank_consistency(last_logits) last_logits_cpu = last_logits.cpu() @@ -248,6 +389,9 @@ def main() -> None: print(f'Logits saved to: {tensor_path.resolve()}', flush=True) print(f'Report saved to: {json_path.resolve()}', flush=True) + if fingerprint_path is not None: + print(f'Rank {rank} weight fingerprints saved to: {fingerprint_path.resolve()}', flush=True) + if dist.is_available() and dist.is_initialized(): dist.barrier() diff --git a/cookbook/client/server/transformer/read_me.md b/cookbook/client/server/transformer/read_me.md index a0c369a4f..cef5000fd 100644 --- a/cookbook/client/server/transformer/read_me.md +++ b/cookbook/client/server/transformer/read_me.md @@ -432,6 +432,69 @@ python3 cookbook/client/server/transformer/diagnostics/compare_dsv4_4layer_logit --output /nas/disk6/ljl/dsv4_ep_diag/results/compare_cli_ep_loop_vs_ep_gmm.json ``` +### 比较 memory_efficient_init 开关后的实际权重 + +最终 logits 不同只能证明初始化路径影响了结果。要定位具体参数，使用相同的 +`ep_loop` 计算路径分别开启和关闭 `memory_efficient_init`，并增加 +`--dump-weight-fingerprints`。诊断只读取每个 rank 的本地参数，并从每个参数均匀 +采样最多 4096 个值，不会将完整权重保存到磁盘，也不会为了展平非连续参数而复制 +完整权重。 + +开启： + +```bash +torchrun --standalone --nproc-per-node=4 \ + cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_cli.py \ + --mode ep_loop \ + --dump-weight-fingerprints \ + --output-dir /nas/disk6/ljl/dsv4_ep_diag/results/memory_efficient_on +``` + +关闭： + +```bash +torchrun --standalone --nproc-per-node=4 \ + cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_cli.py \ + --mode ep_loop \ + --disable-memory-efficient-init \ + --dump-weight-fingerprints \ + --output-dir /nas/disk6/ljl/dsv4_ep_diag/results/memory_efficient_off +``` + +每次运行会按 rank 生成： + +```text +cli_ep_loop_weight_fingerprints_rank0.json +cli_ep_loop_weight_fingerprints_rank1.json +cli_ep_loop_weight_fingerprints_rank2.json +cli_ep_loop_weight_fingerprints_rank3.json +``` + +以关闭状态作为 reference 进行比较： + +```bash +python3 cookbook/client/server/transformer/diagnostics/compare_dsv4_weight_fingerprints.py \ + /nas/disk6/ljl/dsv4_ep_diag/results/memory_efficient_off \ + /nas/disk6/ljl/dsv4_ep_diag/results/memory_efficient_on \ + --mode ep_loop \ + --output /nas/disk6/ljl/dsv4_ep_diag/results/compare_memory_init_weights.json \ + || true +``` + +重点查看比较报告中的以下字段： + +```text +metadata_mismatches shape、stride、dtype 或连续性不同 +value_mismatches 参数采样值不同 +first_difference 第一个不同采样值的本地 flat index 和多维坐标 +rank 不同权重所在的 EP/FSDP rank +name 不同权重的完整参数名 +``` + +如果第一个不一致参数是 `embed_tokens`、attention、norm 或 `lm_head`，优先检查 +普通 FSDP DTensor 重建。如果只有 `mlp.experts.gate_up_proj/down_proj` 不一致，优先 +检查 pre-EP state capture 和 `_scatter_ep_expert_tensor()` 的专家切分。 + 还可以将 `cli__last_logits.pt` 与 C/S 产生的 `_last_logits.pt` 比较。CLI 一致而 C/S 不一致时，重点排查 Ray/服务端权重 分发；CLI 本身已经不一致时，重点排查 FSDP/EP/GMM。 @@ -809,7 +872,8 @@ python3 -m py_compile \ cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_ray.py \ cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_cli.py \ cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits.py \ - cookbook/client/server/transformer/diagnostics/compare_dsv4_4layer_logits.py + cookbook/client/server/transformer/diagnostics/compare_dsv4_4layer_logits.py \ + cookbook/client/server/transformer/diagnostics/compare_dsv4_weight_fingerprints.py bash -n \ cookbook/client/server/transformer/run_dsv4_0731_npu_2node_2npu.sh \ From cc650b951efa62ccaf4e03e648ab5fac98aff718 Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Thu, 3 Sep 2026 16:38:36 +0800 Subject: [PATCH 27/28] group mm --- .../compare_dsv4_weight_fingerprints.py | 86 ++++++++---- .../probe_dsv4_4layer_logits_cli.py | 126 +++++++++++++++--- .../model/transformers/moe/expert_parallel.py | 50 +++++-- .../transformers/strategy/native_fsdp.py | 54 ++++++-- .../model/transformers/transformers.py | 3 + src/twinkle/utils/torch_utils.py | 22 +++ .../test_ep_multi_lora_target_parameters.py | 32 +++++ tests/utils/test_utils.py | 30 +++++ 8 files changed, 335 insertions(+), 68 deletions(-) diff --git a/cookbook/client/server/transformer/diagnostics/compare_dsv4_weight_fingerprints.py b/cookbook/client/server/transformer/diagnostics/compare_dsv4_weight_fingerprints.py index 093f254eb..dabf1172f 100755 --- a/cookbook/client/server/transformer/diagnostics/compare_dsv4_weight_fingerprints.py +++ b/cookbook/client/server/transformer/diagnostics/compare_dsv4_weight_fingerprints.py @@ -17,7 +17,12 @@ 'dtype', 'local_numel', 'local_is_contiguous', + 'local_is_meta', 'sample_count', + 'is_dtensor', + 'placements', + 'device_mesh_shape', + 'device_mesh_dim_names', ) @@ -46,8 +51,12 @@ def load_rank_reports(directory: Path, mode: str) -> dict[int, dict[str, Any]]: return reports -def parameter_map(report: dict[str, Any]) -> dict[str, dict[str, Any]]: - return {record['name']: record for record in report['parameters']} +def tensor_map(report: dict[str, Any]) -> dict[tuple[str, str], dict[str, Any]]: + # Reports written before buffers were added contain only ``parameters``. + records = report.get('tensors') + if records is None: + records = [dict(record, kind='parameter') for record in report['parameters']] + return {(record.get('kind', 'parameter'), record['name']): record for record in records} def unravel_index(index: int, shape: list[int]) -> list[int]: @@ -116,8 +125,9 @@ def main() -> None: value_mismatches = [] missing_parameters = [] unexpected_parameters = [] - compared_parameters = 0 - exact_parameters = 0 + compared_tensors = 0 + exact_tensors = 0 + full_hash_comparisons = 0 for rank in sorted(reference_ranks): reference_report = reference_reports[rank] @@ -125,39 +135,66 @@ def main() -> None: if reference_report['sample_algorithm'] != candidate_report['sample_algorithm']: raise RuntimeError(f'Rank {rank}: sample algorithms differ') - reference_parameters = parameter_map(reference_report) - candidate_parameters = parameter_map(candidate_report) - reference_names = set(reference_parameters) - candidate_names = set(candidate_parameters) - for name in sorted(reference_names - candidate_names): - missing_parameters.append({'rank': rank, 'name': name}) - for name in sorted(candidate_names - reference_names): - unexpected_parameters.append({'rank': rank, 'name': name}) - - for name in sorted(reference_names & candidate_names): - compared_parameters += 1 - reference = reference_parameters[name] - candidate = candidate_parameters[name] + reference_tensors = tensor_map(reference_report) + candidate_tensors = tensor_map(candidate_report) + reference_names = set(reference_tensors) + candidate_names = set(candidate_tensors) + for kind, name in sorted(reference_names - candidate_names): + missing_parameters.append({'rank': rank, 'kind': kind, 'name': name}) + for kind, name in sorted(candidate_names - reference_names): + unexpected_parameters.append({'rank': rank, 'kind': kind, 'name': name}) + + for kind, name in sorted(reference_names & candidate_names): + compared_tensors += 1 + reference = reference_tensors[(kind, name)] + candidate = candidate_tensors[(kind, name)] changed_metadata = { - field: {'reference': reference[field], 'candidate': candidate[field]} + field: {'reference': reference.get(field), 'candidate': candidate.get(field)} for field in METADATA_FIELDS - if reference[field] != candidate[field] + if reference.get(field) != candidate.get(field) } if changed_metadata: metadata_mismatches.append({ 'rank': rank, + 'kind': kind, 'name': name, 'fields': changed_metadata, }) continue - if reference['sample_sha256'] == candidate['sample_sha256']: - exact_parameters += 1 + reference_full_sha = reference.get('full_sha256') + candidate_full_sha = candidate.get('full_sha256') + if bool(reference_full_sha) != bool(candidate_full_sha): + metadata_mismatches.append({ + 'rank': rank, + 'kind': kind, + 'name': name, + 'fields': { + 'full_sha256_available': { + 'reference': bool(reference_full_sha), + 'candidate': bool(candidate_full_sha), + } + }, + }) + continue + if reference_full_sha: + full_hash_comparisons += 1 + values_match = reference_full_sha == candidate_full_sha + comparison_basis = 'full_sha256' + else: + values_match = reference['sample_sha256'] == candidate['sample_sha256'] + comparison_basis = 'sample_sha256' + if values_match: + exact_tensors += 1 continue mismatch = { 'rank': rank, + 'kind': kind, 'name': name, + 'comparison_basis': comparison_basis, 'reference_sha256': reference['sample_sha256'], 'candidate_sha256': candidate['sample_sha256'], + 'reference_full_sha256': reference_full_sha, + 'candidate_full_sha256': candidate_full_sha, } mismatch.update(describe_value_mismatch(reference, candidate)) value_mismatches.append(mismatch) @@ -175,9 +212,10 @@ def main() -> None: 'reference_memory_efficient_init': reference_reports[min(reference_ranks)]['memory_efficient_init'], 'candidate_memory_efficient_init': candidate_reports[min(candidate_ranks)]['memory_efficient_init'], 'ranks': sorted(reference_ranks), - 'compared_parameters': compared_parameters, - 'exact_sampled_parameters': exact_parameters, - 'exact_sampled_fraction': exact_parameters / compared_parameters if compared_parameters else 0.0, + 'compared_tensors': compared_tensors, + 'full_hash_comparisons': full_hash_comparisons, + 'exact_tensors': exact_tensors, + 'exact_tensor_fraction': exact_tensors / compared_tensors if compared_tensors else 0.0, 'metadata_mismatch_count': len(metadata_mismatches), 'value_mismatch_count': len(value_mismatches), 'missing_parameter_count': len(missing_parameters), diff --git a/cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_cli.py b/cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_cli.py index 44684dac5..3441aea27 100644 --- a/cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_cli.py +++ b/cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_cli.py @@ -49,13 +49,24 @@ def parse_args() -> argparse.Namespace: parser.add_argument( '--dump-weight-fingerprints', action='store_true', - help='Save sampled local base-parameter fingerprints for every rank after FSDP/EP wrapping.', + help='Save local parameter and buffer fingerprints for every rank after FSDP/EP wrapping.', ) parser.add_argument( '--weight-fingerprint-samples', type=int, default=4096, - help='Maximum number of evenly spaced values sampled from each local parameter.', + help='Maximum number of evenly spaced values sampled from each local tensor.', + ) + parser.add_argument( + '--weight-fingerprint-full-hash', + action='store_true', + help='Also calculate an exact SHA256 over every local parameter and buffer in bounded CPU chunks.', + ) + parser.add_argument( + '--weight-fingerprint-chunk-mib', + type=int, + default=64, + help='Maximum target CPU chunk size used by --weight-fingerprint-full-hash.', ) return parser.parse_args() @@ -130,6 +141,28 @@ def _parameter_local_tensor(parameter): return value +def _tensor_distribution_metadata(tensor) -> dict[str, Any]: + from torch.distributed.tensor import DTensor + + value = tensor.detach() + if not isinstance(value, DTensor): + return { + 'is_dtensor': False, + 'placements': None, + 'device_mesh_shape': None, + 'device_mesh_dim_names': None, + } + mesh = value.device_mesh + mesh_tensor = getattr(mesh, 'mesh', None) + mesh_dim_names = getattr(mesh, 'mesh_dim_names', None) + return { + 'is_dtensor': True, + 'placements': [repr(placement) for placement in value.placements], + 'device_mesh_shape': list(mesh_tensor.shape) if mesh_tensor is not None else None, + 'device_mesh_dim_names': list(mesh_dim_names) if mesh_dim_names is not None else None, + } + + def _sample_tensor_at_flat_positions(tensor, positions: list[int]): """Sample logical flat positions without flattening/copying the full tensor.""" import torch @@ -151,12 +184,45 @@ def _sample_tensor_at_flat_positions(tensor, positions: list[int]): return tensor[coordinates] -def _sample_parameter(parameter, sample_limit: int) -> dict[str, Any]: +def _full_tensor_sha256(tensor, chunk_mib: int) -> str: + """Hash a logical tensor in bounded chunks without copying it all to CPU.""" + import torch + + digest = hashlib.sha256() + if tensor.numel() == 0: + return digest.hexdigest() + if tensor.ndim == 0: + cpu_chunk = tensor.detach().cpu().contiguous().reshape(1) + digest.update(cpu_chunk.view(torch.uint8).numpy().tobytes()) + return digest.hexdigest() + + row_numel = 1 + for size in tensor.shape[1:]: + row_numel *= int(size) + row_bytes = max(1, row_numel * tensor.element_size()) + rows_per_chunk = max(1, (chunk_mib * 1024 * 1024) // row_bytes) + for start in range(0, int(tensor.shape[0]), rows_per_chunk): + length = min(rows_per_chunk, int(tensor.shape[0]) - start) + cpu_chunk = tensor.narrow(0, start, length).detach().cpu().contiguous() + digest.update(cpu_chunk.view(torch.uint8).numpy().tobytes()) + return digest.hexdigest() + + +def _fingerprint_tensor(tensor, sample_limit: int, *, full_hash: bool, chunk_mib: int) -> dict[str, Any]: import torch - local = _parameter_local_tensor(parameter) - positions = _sample_positions(local.numel(), sample_limit) - if positions: + local = _parameter_local_tensor(tensor) + local_is_meta = bool(getattr(local, 'is_meta', False)) + positions = [] if local_is_meta else _sample_positions(local.numel(), sample_limit) + if local_is_meta: + sample_sha256 = None + sample_values = [] + sample_finite = None + sample_min = None + sample_max = None + sample_sum = None + sample_abs_sum = None + elif positions: sampled = _sample_tensor_at_flat_positions(local, positions).detach().cpu().contiguous() # Hash the original dtype bytes. Converting BF16 directly to NumPy is # not supported by every NumPy version, whereas a byte view is. @@ -177,14 +243,15 @@ def _sample_parameter(parameter, sample_limit: int) -> dict[str, Any]: sample_sum = 0.0 sample_abs_sum = 0.0 - return { - 'global_shape': list(parameter.shape), - 'global_stride': list(parameter.stride()), + result = { + 'global_shape': list(tensor.shape), + 'global_stride': list(tensor.stride()), 'local_shape': list(local.shape), 'local_stride': list(local.stride()), 'dtype': str(local.dtype), 'local_numel': local.numel(), 'local_is_contiguous': local.is_contiguous(), + 'local_is_meta': local_is_meta, 'sample_count': len(positions), 'sample_sha256': sample_sha256, 'sample_finite': sample_finite, @@ -193,7 +260,10 @@ def _sample_parameter(parameter, sample_limit: int) -> dict[str, Any]: 'sample_min': sample_min, 'sample_max': sample_max, 'sample_values': sample_values, + 'full_sha256': _full_tensor_sha256(local, chunk_mib) if full_hash and not local_is_meta else None, } + result.update(_tensor_distribution_metadata(tensor)) + return result def _write_weight_fingerprints(model, args: argparse.Namespace, *, rank: int, world_size: int) -> Path: @@ -201,16 +271,28 @@ def _write_weight_fingerprints(model, args: argparse.Namespace, *, rank: int, wo import torch inner_model = getattr(model, 'model', model) - parameters = [] + tensors = [] with torch.no_grad(): for name, parameter in inner_model.named_parameters(): - # The probe disables LoRA during forward. Adapter slots are seeded - # separately and are not part of the base-weight loading question. - if 'lora_' in name: - continue - record = {'name': name} - record.update(_sample_parameter(parameter, args.weight_fingerprint_samples)) - parameters.append(record) + record = {'kind': 'parameter', 'name': name} + record.update( + _fingerprint_tensor( + parameter, + args.weight_fingerprint_samples, + full_hash=args.weight_fingerprint_full_hash, + chunk_mib=args.weight_fingerprint_chunk_mib, + )) + tensors.append(record) + for name, buffer in inner_model.named_buffers(): + record = {'kind': 'buffer', 'name': name} + record.update( + _fingerprint_tensor( + buffer, + args.weight_fingerprint_samples, + full_hash=args.weight_fingerprint_full_hash, + chunk_mib=args.weight_fingerprint_chunk_mib, + )) + tensors.append(record) payload = { 'execution': 'cli', @@ -221,8 +303,12 @@ def _write_weight_fingerprints(model, args: argparse.Namespace, *, rank: int, wo 'memory_efficient_init': not args.disable_memory_efficient_init, 'sample_algorithm': 'flat_evenly_spaced_integer_v1', 'sample_limit_per_parameter': args.weight_fingerprint_samples, - 'parameter_count': len(parameters), - 'parameters': parameters, + 'full_hash': args.weight_fingerprint_full_hash, + 'full_hash_chunk_mib': args.weight_fingerprint_chunk_mib, + 'parameter_count': sum(record['kind'] == 'parameter' for record in tensors), + 'buffer_count': sum(record['kind'] == 'buffer' for record in tensors), + 'tensor_count': len(tensors), + 'tensors': tensors, } args.output_dir.mkdir(parents=True, exist_ok=True) output = args.output_dir / f'cli_{args.mode}_weight_fingerprints_rank{rank}.json' @@ -240,6 +326,8 @@ def main() -> None: raise SystemExit('--input-ids must contain at least one token ID') if args.weight_fingerprint_samples <= 0: raise SystemExit('--weight-fingerprint-samples must be positive') + if args.weight_fingerprint_chunk_mib <= 0: + raise SystemExit('--weight-fingerprint-chunk-mib must be positive') world_size = int(os.environ.get('WORLD_SIZE', '1')) rank = int(os.environ.get('RANK', '0')) diff --git a/src/twinkle/model/transformers/moe/expert_parallel.py b/src/twinkle/model/transformers/moe/expert_parallel.py index 867e10e89..1ec1f3361 100644 --- a/src/twinkle/model/transformers/moe/expert_parallel.py +++ b/src/twinkle/model/transformers/moe/expert_parallel.py @@ -55,8 +55,15 @@ def apply_expert_parallel( device_mesh: DeviceMesh, config: dict[str, Any] | None = None, ep_fsdp_device_mesh: torch.distributed.DeviceMesh | None = None, + clone_tensor_expert_weights: bool = True, ) -> list[ExpertShardingSpec]: - """Apply expert parallelism to all MoE blocks in the model.""" + """Apply expert parallelism to all MoE blocks in the model. + + ``clone_tensor_expert_weights=False`` is only intended for the native-FSDP + rank-0 broadcast path. That path retains the complete immutable state dict, + so temporary local expert parameters may be views of the retained storage + until the model is converted to meta and materialized by FSDP. + """ cfg = _merge_config(config) # EP info comes from the separate ep_fsdp_device_mesh, not from main mesh @@ -78,7 +85,13 @@ def apply_expert_parallel( specs = [] for block_name, block in find_moe_blocks_with_names(model): - spec = shard_experts(block, ep_world_size, ep_rank, cfg) + spec = shard_experts( + block, + ep_world_size, + ep_rank, + cfg, + clone_tensor_expert_weights=clone_tensor_expert_weights, + ) patch_forward(block, ep_group, ep_world_size, cfg, block_name) specs.append(spec) @@ -119,6 +132,8 @@ def shard_experts( ep_world_size: int, ep_rank: int, cfg: ExpertParallelConfig, + *, + clone_tensor_expert_weights: bool = True, ) -> ExpertShardingSpec: """Shard experts in a MoE block across EP ranks. @@ -127,6 +142,9 @@ def shard_experts( ep_world_size: The world size for expert parallelism. ep_rank: The current rank in the EP group. cfg: Expert parallel configuration. + clone_tensor_expert_weights: Clone local expert slices for long-lived + training storage. The memory-efficient native-FSDP bootstrap may + disable this because it retains the complete immutable source. Returns an ExpertShardingSpec describing the sharding. """ @@ -144,7 +162,12 @@ def shard_experts( block.experts = local_experts is_tensor_experts = False else: - _shard_tensor_experts(block.experts, local_start, local_end) + _shard_tensor_experts( + block.experts, + local_start, + local_end, + clone=clone_tensor_expert_weights, + ) is_tensor_experts = True block._ep_num_experts = num_experts @@ -487,11 +510,16 @@ def _is_moe_experts(experts: Any) -> bool: return False -def _shard_tensor_experts(experts: nn.Module, start: int, end: int) -> None: - experts.gate_up_proj = nn.Parameter( - experts.gate_up_proj.data[start:end].clone(), requires_grad=experts.gate_up_proj.requires_grad) - experts.down_proj = nn.Parameter( - experts.down_proj.data[start:end].clone(), requires_grad=experts.down_proj.requires_grad) +def _shard_tensor_experts(experts: nn.Module, start: int, end: int, *, clone: bool = True) -> None: + + def _slice_parameter(parameter: nn.Parameter) -> nn.Parameter: + local_tensor = parameter.data[start:end] + if clone: + local_tensor = local_tensor.clone() + return nn.Parameter(local_tensor, requires_grad=parameter.requires_grad) + + experts.gate_up_proj = _slice_parameter(experts.gate_up_proj) + experts.down_proj = _slice_parameter(experts.down_proj) if hasattr(experts, 'num_experts'): experts.num_experts = end - start @@ -500,11 +528,9 @@ def _shard_tensor_experts(experts: nn.Module, start: int, end: int) -> None: if not isinstance(target_param_wrapper, TargetParameterLoraWrapper): continue for tenant_name, tenant_tensor in target_param_wrapper.lora_A.items(): - target_param_wrapper.lora_A[tenant_name] = nn.Parameter( - tenant_tensor.data[start:end].clone(), requires_grad=tenant_tensor.requires_grad) + target_param_wrapper.lora_A[tenant_name] = _slice_parameter(tenant_tensor) for tenant_name, tenant_tensor in target_param_wrapper.lora_B.items(): - target_param_wrapper.lora_B[tenant_name] = nn.Parameter( - tenant_tensor.data[start:end].clone(), requires_grad=tenant_tensor.requires_grad) + target_param_wrapper.lora_B[tenant_name] = _slice_parameter(tenant_tensor) def _run_local_experts( diff --git a/src/twinkle/model/transformers/strategy/native_fsdp.py b/src/twinkle/model/transformers/strategy/native_fsdp.py index f1bd6c6ec..049ec2a59 100644 --- a/src/twinkle/model/transformers/strategy/native_fsdp.py +++ b/src/twinkle/model/transformers/strategy/native_fsdp.py @@ -9,7 +9,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Literal, Mapping, Optional, Set from twinkle.utils import DeviceMesh, Platform, get_logger, torch_util -from twinkle.utils.torch_utils import clone_state_dict_to_cpu +from twinkle.utils.torch_utils import snapshot_state_dict_to_cpu from .load_context import fsdp_pretrained_load_context if TYPE_CHECKING: @@ -77,9 +77,21 @@ def capture_pre_ep_state_if_needed(self, model, *, enable_ep: bool) -> None: if local_rank < 0: raise RuntimeError('Native FSDP node-local pre-EP state capture requires node-local rank topology.') is_source_rank = dist.is_available() and dist.is_initialized() and local_rank == 0 - self.set_rank0_pre_ep_full_state_dict(clone_state_dict_to_cpu(model.state_dict()) if is_source_rank else {}) + # Do not clone the complete checkpoint here. The source-rank model + # already owns one full CPU copy, and state_dict() tensors keep that + # storage alive when EP replaces expert parameters and FSDP later moves + # the model skeleton to meta. Cloning this state used to require a + # second full-model-sized allocation on every node-local source rank. + # The retained tensors are immutable until wrap_model() finishes + # broadcasting the rank-local shards. + self.set_rank0_pre_ep_full_state_dict( + snapshot_state_dict_to_cpu(model.state_dict()) if is_source_rank else {}) self._pre_ep_state_captured = True + def can_reuse_pre_ep_tensor_storage(self) -> bool: + """Whether EP may use temporary views into the retained full state.""" + return self._pre_ep_state_captured and self.use_rank0_pretrained_broadcast() + def prepare_adapter_config(self, config_or_dir, *, enable_ep: bool): if not enable_ep: return config_or_dir @@ -163,7 +175,11 @@ def wrap_model(self, model, optimizer=None): original_sd = self._rank0_pre_ep_full_state_dict if is_source_rank else {} else: original_sd = model.state_dict() if is_source_rank else {} - adapter_source_sd = _collect_adapter_source_state(model.state_dict()) + # These tensors are consumed before any training step and only + # need to outlive model.to(meta), so detached CPU views are + # sufficient here as well. Avoid duplicating all preallocated + # Multi-LoRA slots during initialization. + adapter_source_sd = _collect_adapter_source_state(model.state_dict(), clone=False) if is_source_rank: # Multi-LoRA target-parameter slots are installed before EP # and then sharded with their experts. Preserve their full @@ -239,18 +255,30 @@ def wrap_model(self, model, optimizer=None): device_type = self.device_mesh.device_type or 'cuda' expert_shard_specs = _collect_ep_expert_shard_specs(model) if ep_enabled else {} rank_to_ep_rank = _build_rank_to_ep_rank(self.ep_fsdp_device_mesh) if ep_enabled else {} - _broadcast_sharded_state_dict( - model, - original_sd or {}, - device_type=device_type, - expert_shard_specs=expert_shard_specs, - rank_to_ep_rank=rank_to_ep_rank, - adapter_source_sd=adapter_source_sd, - adapter_full_sd=adapter_full_sd, - ) - self._adapter_full_state_dict = None + try: + _broadcast_sharded_state_dict( + model, + original_sd or {}, + device_type=device_type, + expert_shard_specs=expert_shard_specs, + rank_to_ep_rank=rank_to_ep_rank, + adapter_source_sd=adapter_source_sd, + adapter_full_sd=adapter_full_sd, + ) + finally: + # The materialized FSDP/EP shards now own their device + # storage. Release all full CPU checkpoint references even + # when loading fails, otherwise a long-lived Ray actor can + # retain an entire model-sized snapshot. + self._rank0_pre_ep_full_state_dict = None + self._adapter_full_state_dict = None + original_sd = None + adapter_source_sd.clear() + adapter_full_sd.clear() target_device = torch.device(device_type) _broadcast_non_persistent_buffers(model, saved_buffers or {}, device=target_device) + if saved_buffers is not None: + saved_buffers.clear() if hasattr(model, 'tie_weights'): model.tie_weights() diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index 34b18ba8c..401f29989 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -447,11 +447,14 @@ def _maybe_apply_expert_parallel(self): self._ensure_optimizer_dp_groups() model = self.strategy.unwrap_model(self.model) ep_fsdp_mesh = getattr(self.strategy, 'ep_fsdp_device_mesh', None) + can_reuse_storage = getattr(self.strategy, 'can_reuse_pre_ep_tensor_storage', None) + reuse_pre_ep_storage = bool(can_reuse_storage()) if callable(can_reuse_storage) else False apply_expert_parallel( model, self.device_mesh, config=self._expert_parallel_config, ep_fsdp_device_mesh=ep_fsdp_mesh, + clone_tensor_expert_weights=not reuse_pre_ep_storage, ) self._expert_parallel_applied = True diff --git a/src/twinkle/utils/torch_utils.py b/src/twinkle/utils/torch_utils.py index 42bdaa0c9..ab1c6c7d5 100644 --- a/src/twinkle/utils/torch_utils.py +++ b/src/twinkle/utils/torch_utils.py @@ -32,6 +32,28 @@ def clone_state_dict_to_cpu(state_dict: Mapping[str, Any]) -> dict: return cloned +def snapshot_state_dict_to_cpu(state_dict: Mapping[str, Any]) -> dict: + """Keep a detached CPU snapshot without duplicating CPU tensor storage. + + ``state_dict()`` returns detached tensor views of module parameters and + persistent buffers. Keeping those views alive is enough to retain their + storage after the module replaces the corresponding parameters (for + example, while EP shards experts or while FSDP moves the module to meta). + + CPU tensors are therefore retained by reference instead of cloned. A + non-CPU tensor still has to be copied to CPU. Callers must treat the + returned tensors as immutable until the snapshot is released. + """ + snapshot = {} + for key, value in state_dict.items(): + if not hasattr(value, 'detach'): + snapshot[key] = value + continue + value = value.detach() + snapshot[key] = value if value.device.type == 'cpu' else value.cpu() + return snapshot + + def pad_sequence_to_length( tensor: 'torch.Tensor', max_seq_len: int, diff --git a/tests/moe/test_ep_multi_lora_target_parameters.py b/tests/moe/test_ep_multi_lora_target_parameters.py index 8012cc87c..ed8ff608f 100644 --- a/tests/moe/test_ep_multi_lora_target_parameters.py +++ b/tests/moe/test_ep_multi_lora_target_parameters.py @@ -153,6 +153,38 @@ def test_ep_shards_target_parameter_lora_slots_on_meta(): assert all(param.shape[0] == 2 and param.is_meta for param in wrapper.lora_B.values()) +def test_ep_can_reuse_retained_full_tensor_storage_during_memory_efficient_init(): + _ensure_dummy_zmq() + from twinkle.model.transformers.moe.expert_parallel import _shard_tensor_experts + + experts = _FakeTensorExperts() + full_gate_up = experts.gate_up_proj + full_down = experts.down_proj + expected_gate_up = full_gate_up[2:4].detach().clone() + expected_down = full_down[2:4].detach().clone() + + _shard_tensor_experts(experts, 2, 4, clone=False) + + assert experts.gate_up_proj.untyped_storage().data_ptr() == full_gate_up.untyped_storage().data_ptr() + assert experts.down_proj.untyped_storage().data_ptr() == full_down.untyped_storage().data_ptr() + assert torch.equal(experts.gate_up_proj, expected_gate_up) + assert torch.equal(experts.down_proj, expected_down) + + +def test_ep_clones_expert_storage_by_default(): + _ensure_dummy_zmq() + from twinkle.model.transformers.moe.expert_parallel import _shard_tensor_experts + + experts = _FakeTensorExperts() + full_gate_up = experts.gate_up_proj + full_down = experts.down_proj + + _shard_tensor_experts(experts, 2, 4) + + assert experts.gate_up_proj.untyped_storage().data_ptr() != full_gate_up.untyped_storage().data_ptr() + assert experts.down_proj.untyped_storage().data_ptr() != full_down.untyped_storage().data_ptr() + + def test_target_parameter_slot_reset_uses_materialized_ep_local_snapshot(): _ensure_dummy_zmq() from peft import LoraConfig diff --git a/tests/utils/test_utils.py b/tests/utils/test_utils.py index 641da8ad1..4a7263c86 100644 --- a/tests/utils/test_utils.py +++ b/tests/utils/test_utils.py @@ -18,6 +18,7 @@ from twinkle.utils.torch_utils import ( clone_state_dict_to_cpu, selective_log_softmax, + snapshot_state_dict_to_cpu, to_device, ) from twinkle.utils.network import find_free_port @@ -83,6 +84,35 @@ def test_preserves_non_tensors(self): assert cloned['name'] == 'model' +class TestSnapshotStateDictToCpu: + + def test_reuses_cpu_tensor_storage(self): + state = {'w': torch.randn(3, 4), 'b': torch.randn(4)} + snapshot = snapshot_state_dict_to_cpu(state) + + assert snapshot['w'].device == torch.device('cpu') + assert snapshot['w'].data_ptr() == state['w'].data_ptr() + assert torch.equal(snapshot['w'], state['w']) + + def test_snapshot_survives_module_meta_conversion(self): + model = torch.nn.Linear(4, 3) + expected = model.weight.detach().clone() + snapshot = snapshot_state_dict_to_cpu(model.state_dict()) + + model.to(torch.device('meta')) + + assert model.weight.is_meta + assert snapshot['weight'].device == torch.device('cpu') + assert torch.equal(snapshot['weight'], expected) + + def test_preserves_non_tensors(self): + state = {'step': 100, 'name': 'model'} + snapshot = snapshot_state_dict_to_cpu(state) + + assert snapshot['step'] == 100 + assert snapshot['name'] == 'model' + + class TestPadSequenceToLength: def test_right_pad(self): From 5411e8a38353f4d60aa66448faac7aba35f1ccae Mon Sep 17 00:00:00 2001 From: meichangsu1 <1484603386@qq.com> Date: Fri, 4 Sep 2026 19:41:48 +0800 Subject: [PATCH 28/28] remove unused file --- .../diagnostics/compare_dsv4_4layer_logits.py | 76 -- .../compare_dsv4_weight_fingerprints.py | 248 ----- .../diagnostics/probe_dsv4_4layer_logits.py | 202 ---- .../probe_dsv4_4layer_logits_cli.py | 488 --------- .../probe_dsv4_4layer_logits_ray.py | 284 ----- .../diagnostics/test_dsv4_npu_gmm_layout.py | 131 --- cookbook/client/server/transformer/read_me.md | 983 ------------------ .../server/transformer/run_dsv4_0731_npu.sh | 67 -- .../run_dsv4_0731_npu_2node_2npu.sh | 170 --- .../run_dsv4_0731_npu_multinode.sh | 2 - .../run_dsv4_4layer_ep_diagnostic.sh | 37 - .../server_config_dsv4_0731_npu.yaml | 116 --- ...erver_config_dsv4_0731_npu_2node_2npu.yaml | 121 --- ...server_config_dsv4_0731_npu_multinode.yaml | 10 +- .../server_config_dsv4_4layer_diag_ep.yaml | 119 --- .../server_config_dsv4_4layer_diag_no_ep.yaml | 119 --- .../transformer/server_dsv4_0731_npu.py | 12 - .../client/twinkle/demo_client_dataset.py | 131 +++ .../client/twinkle/demo_server_dataset.py | 122 +++ .../dsv4_multi_lora_self_cognition_sft.py | 292 +++--- .../client/twinkle/dsv4_multi_lora_sft.py | 203 ---- .../run_dsv4_0731_npu_2node_2npu_client.sh | 33 - .../twinkle/run_dsv4_0731_npu_client.sh | 26 - .../run_dsv4_0731_npu_multinode_client.sh | 47 - scripts/process_self_cognition_dataset.py | 86 ++ src/twinkle/kernel/ops/ep/__init__.py | 14 - src/twinkle/kernel/ops/moe/npu.py | 14 +- src/twinkle/model/multi_lora.py | 9 +- .../model/transformers/moe/expert_parallel.py | 57 - .../transformers/multi_lora_transformers.py | 5 +- .../transformers/strategy/native_fsdp.py | 39 +- tests/kernel/ops/test_ep_dispatch.py | 10 - tests/kernel/ops/test_moe.py | 12 - 33 files changed, 476 insertions(+), 3809 deletions(-) delete mode 100755 cookbook/client/server/transformer/diagnostics/compare_dsv4_4layer_logits.py delete mode 100755 cookbook/client/server/transformer/diagnostics/compare_dsv4_weight_fingerprints.py delete mode 100755 cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits.py delete mode 100644 cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_cli.py delete mode 100644 cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_ray.py delete mode 100755 cookbook/client/server/transformer/diagnostics/test_dsv4_npu_gmm_layout.py delete mode 100644 cookbook/client/server/transformer/read_me.md delete mode 100755 cookbook/client/server/transformer/run_dsv4_0731_npu.sh delete mode 100755 cookbook/client/server/transformer/run_dsv4_0731_npu_2node_2npu.sh delete mode 100755 cookbook/client/server/transformer/run_dsv4_4layer_ep_diagnostic.sh delete mode 100644 cookbook/client/server/transformer/server_config_dsv4_0731_npu.yaml delete mode 100644 cookbook/client/server/transformer/server_config_dsv4_0731_npu_2node_2npu.yaml delete mode 100644 cookbook/client/server/transformer/server_config_dsv4_4layer_diag_ep.yaml delete mode 100644 cookbook/client/server/transformer/server_config_dsv4_4layer_diag_no_ep.yaml delete mode 100644 cookbook/client/server/transformer/server_dsv4_0731_npu.py create mode 100644 cookbook/client/twinkle/demo_client_dataset.py create mode 100644 cookbook/client/twinkle/demo_server_dataset.py delete mode 100644 cookbook/client/twinkle/dsv4_multi_lora_sft.py delete mode 100755 cookbook/client/twinkle/run_dsv4_0731_npu_2node_2npu_client.sh delete mode 100755 cookbook/client/twinkle/run_dsv4_0731_npu_client.sh delete mode 100755 cookbook/client/twinkle/run_dsv4_0731_npu_multinode_client.sh create mode 100644 scripts/process_self_cognition_dataset.py diff --git a/cookbook/client/server/transformer/diagnostics/compare_dsv4_4layer_logits.py b/cookbook/client/server/transformer/diagnostics/compare_dsv4_4layer_logits.py deleted file mode 100755 index 9fce9ef63..000000000 --- a/cookbook/client/server/transformer/diagnostics/compare_dsv4_4layer_logits.py +++ /dev/null @@ -1,76 +0,0 @@ -#!/usr/bin/env python3 -"""Numerically compare two last-token-logit files produced by the DSV4 probe.""" - -from __future__ import annotations - -import argparse -import json -from pathlib import Path - -import torch - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument('reference', type=Path) - parser.add_argument('candidate', type=Path) - parser.add_argument('--rtol', type=float, default=1e-2) - parser.add_argument('--atol', type=float, default=2e-2) - parser.add_argument('--output', type=Path) - return parser.parse_args() - - -def load_logits(path: Path) -> tuple[str, list[int], torch.Tensor]: - payload = torch.load(path, map_location='cpu', weights_only=True) - return ( - str(payload.get('mode', path.stem)), - list(payload['input_ids']), - payload['last_logits'].float(), - ) - - -def main() -> None: - args = parse_args() - reference_mode, reference_ids, reference = load_logits(args.reference) - candidate_mode, candidate_ids, candidate = load_logits(args.candidate) - - if reference_ids != candidate_ids: - raise RuntimeError(f'Input IDs differ: {reference_ids} != {candidate_ids}') - if reference.shape != candidate.shape: - raise RuntimeError(f'Logit shapes differ: {tuple(reference.shape)} != {tuple(candidate.shape)}') - - difference = (reference - candidate).abs() - close = torch.isclose(reference, candidate, rtol=args.rtol, atol=args.atol) - top_k = min(20, reference.numel()) - reference_top = torch.topk(reference, k=top_k).indices.tolist() - candidate_top = torch.topk(candidate, k=top_k).indices.tolist() - - report = { - 'reference': str(args.reference), - 'reference_mode': reference_mode, - 'candidate': str(args.candidate), - 'candidate_mode': candidate_mode, - 'input_ids': reference_ids, - 'shape': list(reference.shape), - 'rtol': args.rtol, - 'atol': args.atol, - 'allclose': bool(close.all().item()), - 'close_fraction': float(close.float().mean().item()), - 'max_abs_diff': float(difference.max().item()), - 'mean_abs_diff': float(difference.mean().item()), - 'reference_top20': reference_top, - 'candidate_top20': candidate_top, - 'top20_overlap': len(set(reference_top) & set(candidate_top)) / top_k, - } - - output = args.output or args.reference.with_name(f'compare_{reference_mode}_vs_{candidate_mode}.json') - output.parent.mkdir(parents=True, exist_ok=True) - output.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding='utf-8') - print(json.dumps(report, ensure_ascii=False, indent=2)) - print(f'Report saved to: {output.resolve()}') - if not report['allclose']: - raise SystemExit(1) - - -if __name__ == '__main__': - main() diff --git a/cookbook/client/server/transformer/diagnostics/compare_dsv4_weight_fingerprints.py b/cookbook/client/server/transformer/diagnostics/compare_dsv4_weight_fingerprints.py deleted file mode 100755 index dabf1172f..000000000 --- a/cookbook/client/server/transformer/diagnostics/compare_dsv4_weight_fingerprints.py +++ /dev/null @@ -1,248 +0,0 @@ -#!/usr/bin/env python3 -"""Compare rank-local parameter fingerprints from two DSV4 CLI probe runs.""" - -from __future__ import annotations - -import argparse -import json -from pathlib import Path -from typing import Any - - -METADATA_FIELDS = ( - 'global_shape', - 'global_stride', - 'local_shape', - 'local_stride', - 'dtype', - 'local_numel', - 'local_is_contiguous', - 'local_is_meta', - 'sample_count', - 'is_dtensor', - 'placements', - 'device_mesh_shape', - 'device_mesh_dim_names', -) - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument('reference_dir', type=Path) - parser.add_argument('candidate_dir', type=Path) - parser.add_argument('--mode', default='ep_loop', choices=('no_ep', 'ep_loop', 'ep_gmm')) - parser.add_argument('--max-mismatches', type=int, default=50) - parser.add_argument('--output', type=Path) - return parser.parse_args() - - -def load_rank_reports(directory: Path, mode: str) -> dict[int, dict[str, Any]]: - pattern = f'cli_{mode}_weight_fingerprints_rank*.json' - reports = {} - for path in sorted(directory.glob(pattern)): - payload = json.loads(path.read_text(encoding='utf-8')) - rank = int(payload['rank']) - if rank in reports: - raise RuntimeError(f'Duplicate rank {rank} in {directory}') - payload['_path'] = str(path) - reports[rank] = payload - if not reports: - raise FileNotFoundError(f'No files matching {directory / pattern}') - return reports - - -def tensor_map(report: dict[str, Any]) -> dict[tuple[str, str], dict[str, Any]]: - # Reports written before buffers were added contain only ``parameters``. - records = report.get('tensors') - if records is None: - records = [dict(record, kind='parameter') for record in report['parameters']] - return {(record.get('kind', 'parameter'), record['name']): record for record in records} - - -def unravel_index(index: int, shape: list[int]) -> list[int]: - coordinate = [] - remaining = index - for size in reversed(shape): - coordinate.append(remaining % size) - remaining //= size - return list(reversed(coordinate)) - - -def sample_positions(numel: int, count: int) -> list[int]: - if numel <= 0 or count <= 0: - return [] - count = min(numel, count) - if count == 1: - return [0] - return [index * (numel - 1) // (count - 1) for index in range(count)] - - -def describe_value_mismatch(reference: dict[str, Any], candidate: dict[str, Any]) -> dict[str, Any]: - reference_values = reference['sample_values'] - candidate_values = candidate['sample_values'] - if len(reference_values) != len(candidate_values): - return { - 'reference_sample_count': len(reference_values), - 'candidate_sample_count': len(candidate_values), - } - - differences = [abs(float(left) - float(right)) for left, right in zip(reference_values, candidate_values)] - differing = [index for index, difference in enumerate(differences) if difference != 0.0] - result = { - 'different_sample_count': len(differing), - 'sample_count': len(differences), - 'max_abs_sample_diff': max(differences, default=0.0), - 'mean_abs_sample_diff': sum(differences) / len(differences) if differences else 0.0, - } - if differing: - sample_index = differing[0] - flat_positions = sample_positions(reference['local_numel'], reference['sample_count']) - flat_index = flat_positions[sample_index] - result['first_difference'] = { - 'sample_index': sample_index, - 'local_flat_index': flat_index, - 'local_coordinate': unravel_index(flat_index, reference['local_shape']), - 'reference': reference_values[sample_index], - 'candidate': candidate_values[sample_index], - } - return result - - -def main() -> None: - args = parse_args() - if args.max_mismatches <= 0: - raise SystemExit('--max-mismatches must be positive') - - reference_reports = load_rank_reports(args.reference_dir, args.mode) - candidate_reports = load_rank_reports(args.candidate_dir, args.mode) - reference_ranks = set(reference_reports) - candidate_ranks = set(candidate_reports) - if reference_ranks != candidate_ranks: - raise RuntimeError( - f'Rank sets differ: reference={sorted(reference_ranks)}, candidate={sorted(candidate_ranks)}') - - metadata_mismatches = [] - value_mismatches = [] - missing_parameters = [] - unexpected_parameters = [] - compared_tensors = 0 - exact_tensors = 0 - full_hash_comparisons = 0 - - for rank in sorted(reference_ranks): - reference_report = reference_reports[rank] - candidate_report = candidate_reports[rank] - if reference_report['sample_algorithm'] != candidate_report['sample_algorithm']: - raise RuntimeError(f'Rank {rank}: sample algorithms differ') - - reference_tensors = tensor_map(reference_report) - candidate_tensors = tensor_map(candidate_report) - reference_names = set(reference_tensors) - candidate_names = set(candidate_tensors) - for kind, name in sorted(reference_names - candidate_names): - missing_parameters.append({'rank': rank, 'kind': kind, 'name': name}) - for kind, name in sorted(candidate_names - reference_names): - unexpected_parameters.append({'rank': rank, 'kind': kind, 'name': name}) - - for kind, name in sorted(reference_names & candidate_names): - compared_tensors += 1 - reference = reference_tensors[(kind, name)] - candidate = candidate_tensors[(kind, name)] - changed_metadata = { - field: {'reference': reference.get(field), 'candidate': candidate.get(field)} - for field in METADATA_FIELDS - if reference.get(field) != candidate.get(field) - } - if changed_metadata: - metadata_mismatches.append({ - 'rank': rank, - 'kind': kind, - 'name': name, - 'fields': changed_metadata, - }) - continue - reference_full_sha = reference.get('full_sha256') - candidate_full_sha = candidate.get('full_sha256') - if bool(reference_full_sha) != bool(candidate_full_sha): - metadata_mismatches.append({ - 'rank': rank, - 'kind': kind, - 'name': name, - 'fields': { - 'full_sha256_available': { - 'reference': bool(reference_full_sha), - 'candidate': bool(candidate_full_sha), - } - }, - }) - continue - if reference_full_sha: - full_hash_comparisons += 1 - values_match = reference_full_sha == candidate_full_sha - comparison_basis = 'full_sha256' - else: - values_match = reference['sample_sha256'] == candidate['sample_sha256'] - comparison_basis = 'sample_sha256' - if values_match: - exact_tensors += 1 - continue - mismatch = { - 'rank': rank, - 'kind': kind, - 'name': name, - 'comparison_basis': comparison_basis, - 'reference_sha256': reference['sample_sha256'], - 'candidate_sha256': candidate['sample_sha256'], - 'reference_full_sha256': reference_full_sha, - 'candidate_full_sha256': candidate_full_sha, - } - mismatch.update(describe_value_mismatch(reference, candidate)) - value_mismatches.append(mismatch) - - total_mismatches = ( - len(metadata_mismatches) - + len(value_mismatches) - + len(missing_parameters) - + len(unexpected_parameters) - ) - report = { - 'mode': args.mode, - 'reference_dir': str(args.reference_dir), - 'candidate_dir': str(args.candidate_dir), - 'reference_memory_efficient_init': reference_reports[min(reference_ranks)]['memory_efficient_init'], - 'candidate_memory_efficient_init': candidate_reports[min(candidate_ranks)]['memory_efficient_init'], - 'ranks': sorted(reference_ranks), - 'compared_tensors': compared_tensors, - 'full_hash_comparisons': full_hash_comparisons, - 'exact_tensors': exact_tensors, - 'exact_tensor_fraction': exact_tensors / compared_tensors if compared_tensors else 0.0, - 'metadata_mismatch_count': len(metadata_mismatches), - 'value_mismatch_count': len(value_mismatches), - 'missing_parameter_count': len(missing_parameters), - 'unexpected_parameter_count': len(unexpected_parameters), - 'passed': total_mismatches == 0, - 'metadata_mismatches': metadata_mismatches[:args.max_mismatches], - 'value_mismatches': value_mismatches[:args.max_mismatches], - 'missing_parameters': missing_parameters[:args.max_mismatches], - 'unexpected_parameters': unexpected_parameters[:args.max_mismatches], - 'mismatch_lists_truncated': any( - len(items) > args.max_mismatches - for items in ( - metadata_mismatches, - value_mismatches, - missing_parameters, - unexpected_parameters, - ) - ), - } - output = args.output or args.candidate_dir / f'compare_{args.mode}_weight_fingerprints.json' - output.parent.mkdir(parents=True, exist_ok=True) - output.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding='utf-8') - print(json.dumps(report, ensure_ascii=False, indent=2)) - print(f'Report saved to: {output.resolve()}') - if not report['passed']: - raise SystemExit(1) - - -if __name__ == '__main__': - main() diff --git a/cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits.py b/cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits.py deleted file mode 100755 index 30575c592..000000000 --- a/cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits.py +++ /dev/null @@ -1,202 +0,0 @@ -#!/usr/bin/env python3 -"""Save last-token logits from a running four-layer Twinkle server.""" - -from __future__ import annotations - -import argparse -import hashlib -import json -import time -from pathlib import Path -from typing import Any - -import torch -from peft import LoraConfig - -from twinkle_client import init_twinkle_client -from twinkle_client.model import MultiLoraTransformersModel - - -DEFAULT_INPUT_IDS = [0, 128803, 2788, 6573, 70979, 36005, 320, 128804, 128821] -TARGET_PARAMETERS = ['mlp.experts.gate_up_proj', 'mlp.experts.down_proj'] - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument('--mode', required=True, choices=('no_ep', 'ep_loop', 'ep_gmm')) - parser.add_argument('--server-url', default='http://127.0.0.1:8000') - parser.add_argument('--server-token', default='EMPTY_TOKEN') - parser.add_argument('--served-model', default='deepseek-v4-0731-local') - parser.add_argument('--output-dir', type=Path, default=Path('output/dsv4_ep_diag')) - parser.add_argument('--input-ids', default=','.join(str(item) for item in DEFAULT_INPUT_IDS)) - parser.add_argument( - '--capacity-wait-seconds', - type=float, - default=90.0, - help='Wait this long for a previous diagnostic adapter to expire (default: 90).', - ) - parser.add_argument( - '--capacity-poll-seconds', - type=float, - default=5.0, - help='Seconds between LoRA-capacity checks (default: 5).', - ) - return parser.parse_args() - - -def _extract_logits(result: Any) -> torch.Tensor: - if hasattr(result, 'model_dump'): - result = result.model_dump() - if isinstance(result, list) and len(result) == 1 and isinstance(result[0], dict): - result = result[0] - if not isinstance(result, dict) or result.get('logits') is None: - raise RuntimeError(f'forward_only did not return logits; result type={type(result).__name__}') - - logits = torch.as_tensor(result['logits'], dtype=torch.float32) - original_shape = tuple(logits.shape) - while logits.ndim > 3: - logits = logits[0] - if logits.ndim == 3: - logits = logits[0, -1] - elif logits.ndim == 2: - logits = logits[-1] - elif logits.ndim != 1: - raise RuntimeError(f'Unsupported logits shape: {original_shape}') - if logits.numel() < 1000: - raise RuntimeError(f'Last-token logits look too small: shape={tuple(logits.shape)}') - return logits.contiguous() - - -def _tensor_sha256(tensor: torch.Tensor) -> str: - values = tensor.detach().cpu().contiguous().numpy().tobytes() - return hashlib.sha256(values).hexdigest() - - -def _wait_for_free_lora(client: Any, timeout: float, poll_interval: float) -> Any: - """Wait for the prior probe's session-bound adapter to expire.""" - if timeout < 0: - raise ValueError('--capacity-wait-seconds must be non-negative') - if poll_interval <= 0: - raise ValueError('--capacity-poll-seconds must be positive') - - deadline = time.monotonic() + timeout - while True: - capacity = client.get_capacity_info() - if capacity.free_loras >= 1: - return capacity - - remaining = deadline - time.monotonic() - if remaining <= 0: - if capacity.max_loras == 0: - raise RuntimeError( - 'Timed out waiting for ModelManagement capacity: ' - 'max_loras=0, used_loras=0. No model replica registered its ' - 'capacity; inspect the ModelManagement startup and replica ' - 'registration logs.') - raise RuntimeError( - 'Timed out waiting for a free LoRA slot: ' - f'max_loras={capacity.max_loras}, used_loras={capacity.used_loras}, ' - f'free_loras={capacity.free_loras}. The previous diagnostic adapter ' - 'did not expire; inspect the ModelManagement cleanup logs or restart ' - 'the diagnostic server.') - - sleep_seconds = min(poll_interval, remaining) - print( - 'Waiting for a free LoRA slot: ' - f'used={capacity.used_loras}/{capacity.max_loras}, ' - f'remaining={remaining:.1f}s', - flush=True, - ) - time.sleep(sleep_seconds) - - -def main() -> None: - args = parse_args() - input_ids = [int(item.strip()) for item in args.input_ids.split(',') if item.strip()] - if not input_ids: - raise SystemExit('--input-ids must contain at least one token ID') - - client = init_twinkle_client( - base_url=args.server_url, - api_key=args.server_token, - session_heartbeat_interval=10, - ) - try: - # ModelManagement registers its capacity lazily on the first model - # request. ``MultiLoraTransformersModel.__init__`` calls /create, - # ensuring get_capacity_info() reports 0/max_loras instead of 0/0. - model = MultiLoraTransformersModel(model_id=args.served_model) - _wait_for_free_lora( - client, - timeout=args.capacity_wait_seconds, - poll_interval=args.capacity_poll_seconds, - ) - - model.add_adapter_to_model( - f'dsv4_diag_{args.mode}', - LoraConfig( - r=8, - lora_alpha=32, - lora_dropout=0.0, - target_modules=None, - target_parameters=TARGET_PARAMETERS, - bias='none', - ), - gradient_accumulation_steps=1, - ) - model.set_processor('InputProcessor', padding_side='left', padding_free=False) - - raw_input = { - 'input_ids': input_ids, - 'attention_mask': [1] * len(input_ids), - 'position_ids': list(range(len(input_ids))), - } - response = model.forward_only( - inputs=[raw_input], - disable_lora=True, - return_logits=True, - ) - last_logits = _extract_logits(response.result) - finally: - client.close() - - finite = torch.isfinite(last_logits) - top_values, top_indices = torch.topk(last_logits, k=min(20, last_logits.numel())) - report = { - 'mode': args.mode, - 'server_url': args.server_url, - 'served_model': args.served_model, - 'input_ids': input_ids, - 'last_logits_shape': list(last_logits.shape), - 'dtype_saved': str(last_logits.dtype), - 'sha256': _tensor_sha256(last_logits), - 'finite': bool(finite.all().item()), - 'nan_count': int(torch.isnan(last_logits).sum().item()), - 'inf_count': int(torch.isinf(last_logits).sum().item()), - 'sum': last_logits.sum().item(), - 'abs_sum': last_logits.abs().sum().item(), - 'min': last_logits.min().item(), - 'max': last_logits.max().item(), - 'top_token_ids': top_indices.tolist(), - 'top_logits': top_values.tolist(), - } - - args.output_dir.mkdir(parents=True, exist_ok=True) - tensor_path = args.output_dir / f'{args.mode}_last_logits.pt' - json_path = args.output_dir / f'{args.mode}_last_logits.json' - torch.save( - { - 'mode': args.mode, - 'input_ids': input_ids, - 'last_logits': last_logits, - }, - tensor_path, - ) - json_path.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding='utf-8') - print(json.dumps(report, ensure_ascii=False, indent=2)) - print(f'Logits saved to: {tensor_path.resolve()}') - print(f'Report saved to: {json_path.resolve()}') - - -if __name__ == '__main__': - main() diff --git a/cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_cli.py b/cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_cli.py deleted file mode 100644 index 3441aea27..000000000 --- a/cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_cli.py +++ /dev/null @@ -1,488 +0,0 @@ -#!/usr/bin/env python3 -"""Run a four-layer DeepSeek-V4 logits probe directly with torchrun. - -This is the local/CLI counterpart of ``probe_dsv4_4layer_logits.py``. It -constructs ``MultiLoraTransformersModel`` in every torchrun process and -therefore does not use GatewayServer, Ray Serve, sessions, tenants, or -server-side LoRA capacity. A disabled diagnostic adapter is installed to -match the production initialization path. -""" - -from __future__ import annotations - -import argparse -import hashlib -import json -import os -from collections.abc import Mapping -from pathlib import Path -from typing import Any - - -DEFAULT_INPUT_IDS = [0, 128803, 2788, 6573, 70979, 36005, 320, 128804, 128821] -TARGET_PARAMETERS = ['mlp.experts.gate_up_proj', 'mlp.experts.down_proj'] - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument('--mode', required=True, choices=('no_ep', 'ep_loop', 'ep_gmm')) - parser.add_argument( - '--model-id', - default=os.environ.get('DSV4_MODEL_ID'), - help='Four-layer BF16 checkpoint. Defaults to DSV4_MODEL_ID.', - ) - parser.add_argument('--output-dir', type=Path, default=Path('/nas/disk6/ljl/dsv4_ep_diag/results')) - parser.add_argument('--input-ids', default=','.join(str(item) for item in DEFAULT_INPUT_IDS)) - parser.add_argument('--seed', type=int, default=42) - parser.add_argument('--mixed-precision', choices=('no', 'fp16', 'bf16'), default='bf16') - parser.add_argument( - '--ep-size', - type=int, - default=None, - help='EP size for ep_loop/ep_gmm. Defaults to the torchrun world size.', - ) - parser.add_argument( - '--disable-memory-efficient-init', - action='store_true', - help='Load the complete checkpoint in every process before FSDP wrapping.', - ) - parser.add_argument( - '--dump-weight-fingerprints', - action='store_true', - help='Save local parameter and buffer fingerprints for every rank after FSDP/EP wrapping.', - ) - parser.add_argument( - '--weight-fingerprint-samples', - type=int, - default=4096, - help='Maximum number of evenly spaced values sampled from each local tensor.', - ) - parser.add_argument( - '--weight-fingerprint-full-hash', - action='store_true', - help='Also calculate an exact SHA256 over every local parameter and buffer in bounded CPU chunks.', - ) - parser.add_argument( - '--weight-fingerprint-chunk-mib', - type=int, - default=64, - help='Maximum target CPU chunk size used by --weight-fingerprint-full-hash.', - ) - return parser.parse_args() - - -def _extract_logits(result: Any): - import torch - - if hasattr(result, 'to_dict'): - result = result.to_dict() - elif hasattr(result, 'model_dump'): - result = result.model_dump() - if isinstance(result, list) and len(result) == 1 and isinstance(result[0], Mapping): - result = result[0] - if not isinstance(result, Mapping) or result.get('logits') is None: - raise RuntimeError(f'forward_only did not return logits; result type={type(result).__name__}') - - logits = result['logits'] - if not isinstance(logits, torch.Tensor): - logits = torch.as_tensor(logits) - logits = logits.detach().to(dtype=torch.float32) - original_shape = tuple(logits.shape) - while logits.ndim > 3: - logits = logits[0] - if logits.ndim == 3: - logits = logits[0, -1] - elif logits.ndim == 2: - logits = logits[-1] - elif logits.ndim != 1: - raise RuntimeError(f'Unsupported logits shape: {original_shape}') - if logits.numel() < 1000: - raise RuntimeError(f'Last-token logits look too small: shape={tuple(logits.shape)}') - return logits.contiguous() - - -def _tensor_sha256(tensor) -> str: - values = tensor.detach().cpu().contiguous().numpy().tobytes() - return hashlib.sha256(values).hexdigest() - - -def _check_rank_consistency(last_logits): - """Return the largest rank-to-rank difference without gathering Python objects.""" - import torch - import torch.distributed as dist - - if not dist.is_available() or not dist.is_initialized() or dist.get_world_size() == 1: - return 0.0 - - rank0_logits = last_logits.clone() - dist.broadcast(rank0_logits, src=0) - max_diff = (last_logits - rank0_logits).abs().max() - dist.all_reduce(max_diff, op=dist.ReduceOp.MAX) - return float(max_diff.item()) - - -def _sample_positions(numel: int, count: int) -> list[int]: - """Return deterministic, evenly spaced flat indices without float rounding.""" - if numel <= 0 or count <= 0: - return [] - count = min(numel, count) - if count == 1: - return [0] - return [index * (numel - 1) // (count - 1) for index in range(count)] - - -def _parameter_local_tensor(parameter): - """Return the rank-local tensor for both regular Parameters and FSDP2 DTensors.""" - from torch.distributed.tensor import DTensor - - value = parameter.detach() - if isinstance(value, DTensor): - value = value.to_local() - return value - - -def _tensor_distribution_metadata(tensor) -> dict[str, Any]: - from torch.distributed.tensor import DTensor - - value = tensor.detach() - if not isinstance(value, DTensor): - return { - 'is_dtensor': False, - 'placements': None, - 'device_mesh_shape': None, - 'device_mesh_dim_names': None, - } - mesh = value.device_mesh - mesh_tensor = getattr(mesh, 'mesh', None) - mesh_dim_names = getattr(mesh, 'mesh_dim_names', None) - return { - 'is_dtensor': True, - 'placements': [repr(placement) for placement in value.placements], - 'device_mesh_shape': list(mesh_tensor.shape) if mesh_tensor is not None else None, - 'device_mesh_dim_names': list(mesh_dim_names) if mesh_dim_names is not None else None, - } - - -def _sample_tensor_at_flat_positions(tensor, positions: list[int]): - """Sample logical flat positions without flattening/copying the full tensor.""" - import torch - - if not positions: - return torch.empty(0, dtype=tensor.dtype, device=tensor.device) - if tensor.ndim == 0: - return tensor.reshape(1) - - remaining = positions - coordinate_columns = [] - for size in reversed(tensor.shape): - coordinate_columns.append([position % size for position in remaining]) - remaining = [position // size for position in remaining] - coordinates = tuple( - torch.tensor(column, dtype=torch.long, device=tensor.device) - for column in reversed(coordinate_columns) - ) - return tensor[coordinates] - - -def _full_tensor_sha256(tensor, chunk_mib: int) -> str: - """Hash a logical tensor in bounded chunks without copying it all to CPU.""" - import torch - - digest = hashlib.sha256() - if tensor.numel() == 0: - return digest.hexdigest() - if tensor.ndim == 0: - cpu_chunk = tensor.detach().cpu().contiguous().reshape(1) - digest.update(cpu_chunk.view(torch.uint8).numpy().tobytes()) - return digest.hexdigest() - - row_numel = 1 - for size in tensor.shape[1:]: - row_numel *= int(size) - row_bytes = max(1, row_numel * tensor.element_size()) - rows_per_chunk = max(1, (chunk_mib * 1024 * 1024) // row_bytes) - for start in range(0, int(tensor.shape[0]), rows_per_chunk): - length = min(rows_per_chunk, int(tensor.shape[0]) - start) - cpu_chunk = tensor.narrow(0, start, length).detach().cpu().contiguous() - digest.update(cpu_chunk.view(torch.uint8).numpy().tobytes()) - return digest.hexdigest() - - -def _fingerprint_tensor(tensor, sample_limit: int, *, full_hash: bool, chunk_mib: int) -> dict[str, Any]: - import torch - - local = _parameter_local_tensor(tensor) - local_is_meta = bool(getattr(local, 'is_meta', False)) - positions = [] if local_is_meta else _sample_positions(local.numel(), sample_limit) - if local_is_meta: - sample_sha256 = None - sample_values = [] - sample_finite = None - sample_min = None - sample_max = None - sample_sum = None - sample_abs_sum = None - elif positions: - sampled = _sample_tensor_at_flat_positions(local, positions).detach().cpu().contiguous() - # Hash the original dtype bytes. Converting BF16 directly to NumPy is - # not supported by every NumPy version, whereas a byte view is. - sample_sha256 = hashlib.sha256(sampled.view(torch.uint8).numpy().tobytes()).hexdigest() - sampled_float = sampled.float() - sample_values = sampled_float.tolist() - sample_finite = bool(torch.isfinite(sampled_float).all().item()) - sample_min = float(sampled_float.min().item()) - sample_max = float(sampled_float.max().item()) - sample_sum = float(sampled_float.sum().item()) - sample_abs_sum = float(sampled_float.abs().sum().item()) - else: - sample_sha256 = hashlib.sha256(b'').hexdigest() - sample_values = [] - sample_finite = True - sample_min = None - sample_max = None - sample_sum = 0.0 - sample_abs_sum = 0.0 - - result = { - 'global_shape': list(tensor.shape), - 'global_stride': list(tensor.stride()), - 'local_shape': list(local.shape), - 'local_stride': list(local.stride()), - 'dtype': str(local.dtype), - 'local_numel': local.numel(), - 'local_is_contiguous': local.is_contiguous(), - 'local_is_meta': local_is_meta, - 'sample_count': len(positions), - 'sample_sha256': sample_sha256, - 'sample_finite': sample_finite, - 'sample_sum': sample_sum, - 'sample_abs_sum': sample_abs_sum, - 'sample_min': sample_min, - 'sample_max': sample_max, - 'sample_values': sample_values, - 'full_sha256': _full_tensor_sha256(local, chunk_mib) if full_hash and not local_is_meta else None, - } - result.update(_tensor_distribution_metadata(tensor)) - return result - - -def _write_weight_fingerprints(model, args: argparse.Namespace, *, rank: int, world_size: int) -> Path: - """Write rank-local post-wrap base-weight fingerprints to a JSON file.""" - import torch - - inner_model = getattr(model, 'model', model) - tensors = [] - with torch.no_grad(): - for name, parameter in inner_model.named_parameters(): - record = {'kind': 'parameter', 'name': name} - record.update( - _fingerprint_tensor( - parameter, - args.weight_fingerprint_samples, - full_hash=args.weight_fingerprint_full_hash, - chunk_mib=args.weight_fingerprint_chunk_mib, - )) - tensors.append(record) - for name, buffer in inner_model.named_buffers(): - record = {'kind': 'buffer', 'name': name} - record.update( - _fingerprint_tensor( - buffer, - args.weight_fingerprint_samples, - full_hash=args.weight_fingerprint_full_hash, - chunk_mib=args.weight_fingerprint_chunk_mib, - )) - tensors.append(record) - - payload = { - 'execution': 'cli', - 'mode': args.mode, - 'model_id': str(Path(args.model_id).expanduser()), - 'rank': rank, - 'world_size': world_size, - 'memory_efficient_init': not args.disable_memory_efficient_init, - 'sample_algorithm': 'flat_evenly_spaced_integer_v1', - 'sample_limit_per_parameter': args.weight_fingerprint_samples, - 'full_hash': args.weight_fingerprint_full_hash, - 'full_hash_chunk_mib': args.weight_fingerprint_chunk_mib, - 'parameter_count': sum(record['kind'] == 'parameter' for record in tensors), - 'buffer_count': sum(record['kind'] == 'buffer' for record in tensors), - 'tensor_count': len(tensors), - 'tensors': tensors, - } - args.output_dir.mkdir(parents=True, exist_ok=True) - output = args.output_dir / f'cli_{args.mode}_weight_fingerprints_rank{rank}.json' - output.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding='utf-8') - return output - - -def main() -> None: - args = parse_args() - if not args.model_id: - raise SystemExit('Set DSV4_MODEL_ID or pass --model-id.') - - input_ids = [int(item.strip()) for item in args.input_ids.split(',') if item.strip()] - if not input_ids: - raise SystemExit('--input-ids must contain at least one token ID') - if args.weight_fingerprint_samples <= 0: - raise SystemExit('--weight-fingerprint-samples must be positive') - if args.weight_fingerprint_chunk_mib <= 0: - raise SystemExit('--weight-fingerprint-chunk-mib must be positive') - - world_size = int(os.environ.get('WORLD_SIZE', '1')) - rank = int(os.environ.get('RANK', '0')) - enable_ep = args.mode != 'no_ep' - ep_size = args.ep_size if args.ep_size is not None else (world_size if enable_ep else 1) - if enable_ep and world_size <= 1: - raise SystemExit(f'{args.mode} requires at least two torchrun processes; WORLD_SIZE={world_size}.') - if enable_ep and (ep_size <= 1 or world_size % ep_size != 0): - raise SystemExit(f'Invalid EP topology: WORLD_SIZE={world_size}, ep_size={ep_size}.') - if not enable_ep and args.ep_size not in (None, 1): - raise SystemExit('no_ep only supports --ep-size 1.') - - # Kernel selection is read while EP patches are installed. Set it before - # importing Twinkle model modules so CLI and server modes choose the same - # implementation. - os.environ['TWINKLE_EP_FORCE_LOOP'] = '1' if args.mode == 'ep_loop' else '0' - os.environ.setdefault('TWINKLE_EP_DIAGNOSTICS', '1') - os.environ.setdefault('TWINKLE_TRUST_REMOTE_CODE', '1') - - import torch - import torch.distributed as dist - from peft import LoraConfig - from transformers import AutoConfig - - import twinkle - from twinkle import DeviceMesh, Platform - from twinkle.model import MultiLoraTransformersModel - - device_mesh = DeviceMesh.from_sizes( - fsdp_size=world_size, - dp_size=1, - ep_size=ep_size, - device_type=Platform.get_platform().device_prefix(), - ) - twinkle.initialize( - mode='local', - global_device_mesh=device_mesh, - seed=args.seed, - lazy_collect=False, - ) - - config = AutoConfig.from_pretrained(args.model_id, trust_remote_code=True) - adapter_name = f'dsv4_diag_{args.mode}' - model = MultiLoraTransformersModel( - model_id=args.model_id, - config=config, - device_mesh=device_mesh, - strategy='native_fsdp', - mixed_precision=args.mixed_precision, - memory_efficient_init=not args.disable_memory_efficient_init, - max_loras=1, - max_r=8, - max_length=512, - target_modules='all-linear', - fsdp_config={ - 'reshard_after_forward': True, - 'expert_parallel': { - 'enabled': enable_ep, - 'ep_size': ep_size, - 'router_dtype': 'fp32', - 'keep_router_logits': False, - }, - }, - ) - model.add_adapter_to_model( - adapter_name, - LoraConfig( - r=8, - lora_alpha=32, - lora_dropout=0.0, - target_modules=None, - target_parameters=TARGET_PARAMETERS, - bias='none', - ), - gradient_accumulation_steps=1, - ) - model.set_processor( - 'InputProcessor', - adapter_name=adapter_name, - padding_side='left', - padding_free=False, - ) - - raw_input = { - 'input_ids': input_ids, - 'attention_mask': [1] * len(input_ids), - 'position_ids': list(range(len(input_ids))), - } - response = model.forward_only( - inputs=[raw_input], - adapter_name=adapter_name, - disable_lora=True, - return_logits=True, - ) - last_logits = _extract_logits(response) - fingerprint_path = None - if args.dump_weight_fingerprints: - fingerprint_path = _write_weight_fingerprints( - model, - args, - rank=rank, - world_size=world_size, - ) - max_rank_diff = _check_rank_consistency(last_logits) - last_logits_cpu = last_logits.cpu() - - if rank == 0: - finite = torch.isfinite(last_logits_cpu) - top_values, top_indices = torch.topk(last_logits_cpu, k=min(20, last_logits_cpu.numel())) - report = { - 'execution': 'cli', - 'mode': args.mode, - 'model_id': str(Path(args.model_id).expanduser()), - 'world_size': world_size, - 'ep_size': ep_size, - 'memory_efficient_init': not args.disable_memory_efficient_init, - 'input_ids': input_ids, - 'last_logits_shape': list(last_logits_cpu.shape), - 'dtype_saved': str(last_logits_cpu.dtype), - 'sha256': _tensor_sha256(last_logits_cpu), - 'max_rank_diff': max_rank_diff, - 'finite': bool(finite.all().item()), - 'nan_count': int(torch.isnan(last_logits_cpu).sum().item()), - 'inf_count': int(torch.isinf(last_logits_cpu).sum().item()), - 'sum': last_logits_cpu.sum().item(), - 'abs_sum': last_logits_cpu.abs().sum().item(), - 'min': last_logits_cpu.min().item(), - 'max': last_logits_cpu.max().item(), - 'top_token_ids': top_indices.tolist(), - 'top_logits': top_values.tolist(), - } - - args.output_dir.mkdir(parents=True, exist_ok=True) - tensor_path = args.output_dir / f'cli_{args.mode}_last_logits.pt' - json_path = args.output_dir / f'cli_{args.mode}_last_logits.json' - torch.save( - { - 'execution': 'cli', - 'mode': args.mode, - 'input_ids': input_ids, - 'last_logits': last_logits_cpu, - }, - tensor_path, - ) - json_path.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding='utf-8') - print(json.dumps(report, ensure_ascii=False, indent=2), flush=True) - print(f'Logits saved to: {tensor_path.resolve()}', flush=True) - print(f'Report saved to: {json_path.resolve()}', flush=True) - - if fingerprint_path is not None: - print(f'Rank {rank} weight fingerprints saved to: {fingerprint_path.resolve()}', flush=True) - - if dist.is_available() and dist.is_initialized(): - dist.barrier() - - -if __name__ == '__main__': - main() diff --git a/cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_ray.py b/cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_ray.py deleted file mode 100644 index bed3e0214..000000000 --- a/cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_ray.py +++ /dev/null @@ -1,284 +0,0 @@ -#!/usr/bin/env python3 -"""Run the four-layer DeepSeek-V4 logits probe on a multi-node Ray cluster. - -The script is a command-line Ray driver. It creates distributed Twinkle model -actors directly and deliberately does not deploy Ray Serve, GatewayServer, -ModelManagement, sessions, tenants, or server-side LoRA capacity. A disabled -diagnostic adapter is installed inside the model to match the production -MultiLoraTransformersModel initialization path. -""" - -from __future__ import annotations - -import argparse -import hashlib -import json -import os -from collections.abc import Mapping -from pathlib import Path -from typing import Any - - -DEFAULT_INPUT_IDS = [0, 128803, 2788, 6573, 70979, 36005, 320, 128804, 128821] -TARGET_PARAMETERS = ['mlp.experts.gate_up_proj', 'mlp.experts.down_proj'] - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument('--mode', required=True, choices=('no_ep', 'ep_loop', 'ep_gmm')) - parser.add_argument( - '--model-id', - default=os.environ.get('DSV4_MODEL_ID'), - help='Four-layer BF16 checkpoint. Defaults to DSV4_MODEL_ID.', - ) - parser.add_argument('--ray-address', default=os.environ.get('RAY_ADDRESS', 'auto')) - parser.add_argument('--world-size', type=int, default=4) - parser.add_argument('--nproc-per-node', type=int, default=2) - parser.add_argument('--output-dir', type=Path, default=Path('/nas/disk6/ljl/dsv4_ep_diag/results')) - parser.add_argument('--input-ids', default=','.join(str(item) for item in DEFAULT_INPUT_IDS)) - parser.add_argument('--seed', type=int, default=42) - parser.add_argument('--mixed-precision', choices=('no', 'fp16', 'bf16'), default='bf16') - parser.add_argument( - '--ep-size', - type=int, - default=None, - help='EP size for ep_loop/ep_gmm. Defaults to world-size.', - ) - parser.add_argument( - '--disable-memory-efficient-init', - action='store_true', - help='Load the complete checkpoint in every actor before FSDP wrapping.', - ) - return parser.parse_args() - - -def _extract_logits(result: Any): - import torch - - if callable(result): - result = result() - if hasattr(result, 'to_dict'): - result = result.to_dict() - elif hasattr(result, 'model_dump'): - result = result.model_dump() - if isinstance(result, list) and len(result) == 1 and isinstance(result[0], Mapping): - result = result[0] - if not isinstance(result, Mapping) or result.get('logits') is None: - raise RuntimeError(f'forward_only did not return logits; result type={type(result).__name__}') - - logits = result['logits'] - if not isinstance(logits, torch.Tensor): - logits = torch.as_tensor(logits) - logits = logits.detach().cpu().to(dtype=torch.float32) - original_shape = tuple(logits.shape) - while logits.ndim > 3: - logits = logits[0] - if logits.ndim == 3: - logits = logits[0, -1] - elif logits.ndim == 2: - logits = logits[-1] - elif logits.ndim != 1: - raise RuntimeError(f'Unsupported logits shape: {original_shape}') - if logits.numel() < 1000: - raise RuntimeError(f'Last-token logits look too small: shape={tuple(logits.shape)}') - return logits.contiguous() - - -def _tensor_sha256(tensor) -> str: - values = tensor.detach().cpu().contiguous().numpy().tobytes() - return hashlib.sha256(values).hexdigest() - - -def _validate_topology(args: argparse.Namespace, ray) -> None: - required_nodes = args.world_size // args.nproc_per_node - alive_nodes = [node for node in ray.nodes() if node.get('Alive', True)] - npu_nodes = [ - node for node in alive_nodes - if int(node.get('Resources', {}).get('NPU', 0)) >= args.nproc_per_node - ] - total_npus = int(ray.cluster_resources().get('NPU', 0)) - if total_npus < args.world_size or len(npu_nodes) < required_nodes: - raise RuntimeError( - 'Ray cluster does not have the requested NPU topology: ' - f'required world_size={args.world_size}, nodes={required_nodes}, ' - f'nproc_per_node={args.nproc_per_node}; found NPU={total_npus}, ' - f'eligible_nodes={len(npu_nodes)}. Start every Ray node with ' - f'--resources=\'{{"NPU": {args.nproc_per_node}}}\'.') - - -def main() -> None: - args = parse_args() - if not args.model_id: - raise SystemExit('Set DSV4_MODEL_ID or pass --model-id.') - if args.world_size <= 0 or args.nproc_per_node <= 0: - raise SystemExit('--world-size and --nproc-per-node must be positive.') - if args.world_size % args.nproc_per_node != 0: - raise SystemExit('--world-size must be divisible by --nproc-per-node.') - - input_ids = [int(item.strip()) for item in args.input_ids.split(',') if item.strip()] - if not input_ids: - raise SystemExit('--input-ids must contain at least one token ID') - - enable_ep = args.mode != 'no_ep' - ep_size = args.ep_size if args.ep_size is not None else (args.world_size if enable_ep else 1) - if enable_ep and (ep_size <= 1 or args.world_size % ep_size != 0): - raise SystemExit(f'Invalid EP topology: world_size={args.world_size}, ep_size={ep_size}.') - if not enable_ep and args.ep_size not in (None, 1): - raise SystemExit('no_ep only supports --ep-size 1.') - - # RayHelper copies the driver's environment into each model actor. Set EP - # selection before Twinkle creates the placement groups and runtime_envs. - os.environ['TWINKLE_EP_FORCE_LOOP'] = '1' if args.mode == 'ep_loop' else '0' - os.environ.setdefault('TWINKLE_EP_DIAGNOSTICS', '1') - os.environ.setdefault('TWINKLE_TRUST_REMOTE_CODE', '1') - os.environ.setdefault('RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES', '1') - - import ray - import torch - - ray.init(address=args.ray_address, ignore_reinit_error=True) - twinkle_initialized = False - try: - _validate_topology(args, ray) - - from peft import LoraConfig - - import twinkle - from twinkle import DeviceGroup, DeviceMesh - from twinkle.model import MultiLoraTransformersModel - - device_mesh = DeviceMesh.from_sizes( - fsdp_size=args.world_size, - dp_size=1, - ep_size=ep_size, - device_type='npu', - ) - groups = [ - DeviceGroup( - name='model', - ranks=list(range(args.world_size)), - device_type='NPU', - ) - ] - twinkle.initialize( - mode='ray', - nproc_per_node=args.nproc_per_node, - ncpu_proc_per_node=1, - seed=args.seed, - groups=groups, - global_device_mesh=device_mesh, - lazy_collect=False, - ) - twinkle_initialized = True - - # Match the model initialization used by ModelManagement. The - # Multi-LoRA constructor installs its preallocated slots and aligns - # their dtype before native FSDP wrapping. There is no server-side - # session or capacity accounting in this direct-Ray driver. - adapter_name = f'dsv4_diag_{args.mode}' - model = MultiLoraTransformersModel( - model_id=args.model_id, - device_mesh=device_mesh, - remote_group='model', - instance_id=f'dsv4_4layer_ray_{args.mode}', - strategy='native_fsdp', - mixed_precision=args.mixed_precision, - memory_efficient_init=not args.disable_memory_efficient_init, - max_loras=1, - max_r=8, - max_length=512, - target_modules='all-linear', - fsdp_config={ - 'reshard_after_forward': True, - 'expert_parallel': { - 'enabled': enable_ep, - 'ep_size': ep_size, - 'router_dtype': 'fp32', - 'keep_router_logits': False, - }, - }, - ) - model.add_adapter_to_model( - adapter_name, - LoraConfig( - r=8, - lora_alpha=32, - lora_dropout=0.0, - target_modules=None, - target_parameters=TARGET_PARAMETERS, - bias='none', - ), - gradient_accumulation_steps=1, - ) - model.set_processor( - 'InputProcessor', - adapter_name=adapter_name, - padding_side='left', - padding_free=False, - ) - - raw_input = { - 'input_ids': input_ids, - 'attention_mask': [1] * len(input_ids), - 'position_ids': list(range(len(input_ids))), - } - response = model.forward_only( - inputs=[raw_input], - adapter_name=adapter_name, - disable_lora=True, - return_logits=True, - ) - last_logits = _extract_logits(response) - - finite = torch.isfinite(last_logits) - top_values, top_indices = torch.topk(last_logits, k=min(20, last_logits.numel())) - report = { - 'execution': 'ray_cli', - 'mode': args.mode, - 'ray_address': args.ray_address, - 'model_id': args.model_id, - 'world_size': args.world_size, - 'nproc_per_node': args.nproc_per_node, - 'ep_size': ep_size, - 'memory_efficient_init': not args.disable_memory_efficient_init, - 'input_ids': input_ids, - 'last_logits_shape': list(last_logits.shape), - 'dtype_saved': str(last_logits.dtype), - 'sha256': _tensor_sha256(last_logits), - 'finite': bool(finite.all().item()), - 'nan_count': int(torch.isnan(last_logits).sum().item()), - 'inf_count': int(torch.isinf(last_logits).sum().item()), - 'sum': last_logits.sum().item(), - 'abs_sum': last_logits.abs().sum().item(), - 'min': last_logits.min().item(), - 'max': last_logits.max().item(), - 'top_token_ids': top_indices.tolist(), - 'top_logits': top_values.tolist(), - } - - args.output_dir.mkdir(parents=True, exist_ok=True) - tensor_path = args.output_dir / f'ray_{args.mode}_last_logits.pt' - json_path = args.output_dir / f'ray_{args.mode}_last_logits.json' - torch.save( - { - 'execution': 'ray_cli', - 'mode': args.mode, - 'input_ids': input_ids, - 'last_logits': last_logits, - }, - tensor_path, - ) - json_path.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding='utf-8') - print(json.dumps(report, ensure_ascii=False, indent=2), flush=True) - print(f'Logits saved to: {tensor_path.resolve()}', flush=True) - print(f'Report saved to: {json_path.resolve()}', flush=True) - finally: - if twinkle_initialized: - from twinkle.infra._ray import RayHelper - - RayHelper.teardown() - ray.shutdown() - - -if __name__ == '__main__': - main() diff --git a/cookbook/client/server/transformer/diagnostics/test_dsv4_npu_gmm_layout.py b/cookbook/client/server/transformer/diagnostics/test_dsv4_npu_gmm_layout.py deleted file mode 100755 index bd65df18d..000000000 --- a/cookbook/client/server/transformer/diagnostics/test_dsv4_npu_gmm_layout.py +++ /dev/null @@ -1,131 +0,0 @@ -#!/usr/bin/env python3 -"""Compare the DeepSeek-V4 square expert layout on NPU GMM against F.linear.""" - -from __future__ import annotations - -import argparse -import json -from pathlib import Path - -import torch -import torch.nn.functional as F -from torch import nn - -from twinkle.kernel.ops.moe.npu import GmmFunction, _normalize_packed_expert_weights - - -class PackedExperts(nn.Module): - - def __init__(self, gate_up_proj: torch.Tensor, down_proj: torch.Tensor): - super().__init__() - self.gate_up_proj = nn.Parameter(gate_up_proj, requires_grad=False) - self.down_proj = nn.Parameter(down_proj, requires_grad=False) - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument('--device', default='npu:0') - parser.add_argument('--dtype', choices=('float16', 'bfloat16'), default='bfloat16') - parser.add_argument('--atol', type=float, default=2e-2) - parser.add_argument('--rtol', type=float, default=1e-2) - parser.add_argument('--output', type=Path, default=Path('output/dsv4_ep_diag/npu_gmm_layout.json')) - return parser.parse_args() - - -def reference_forward( - inputs: torch.Tensor, - counts: list[int], - gate_up_proj: torch.Tensor, - down_proj: torch.Tensor, -) -> torch.Tensor: - outputs = [] - start = 0 - for expert, count in enumerate(counts): - expert_input = inputs[start:start + count] - gate_up = F.linear(expert_input, gate_up_proj[expert]) - gate, up = gate_up.chunk(2, dim=-1) - outputs.append(F.linear(F.silu(gate) * up, down_proj[expert])) - start += count - return torch.cat(outputs, dim=0) - - -def gmm_forward( - inputs: torch.Tensor, - counts: torch.Tensor, - gate_up_weight: torch.Tensor, - down_weight: torch.Tensor, -) -> torch.Tensor: - import torch_npu - - gate_up = GmmFunction.apply(inputs, counts, gate_up_weight) - activated = torch_npu.npu_swiglu(gate_up, dim=-1) - return GmmFunction.apply(activated, counts, down_weight) - - -def main() -> None: - args = parse_args() - try: - import torch_npu # noqa: F401 - except ImportError as exc: - raise SystemExit('torch_npu is required; run this script in the Ascend container.') from exc - - if not torch.npu.is_available(): - raise SystemExit('torch.npu.is_available() is False') - - device = torch.device(args.device) - dtype = getattr(torch, args.dtype) - torch.npu.set_device(device.index or 0) - torch.manual_seed(20260901) - torch.npu.manual_seed_all(20260901) - - # Preserve the DeepSeek-V4 relation hidden == 2 * intermediate. The gate/up - # matrix is deliberately non-symmetric so an omitted transpose is visible. - experts = 2 - hidden = 64 - intermediate = 32 - token_counts = [8, 8] - inputs = torch.randn(sum(token_counts), hidden, device=device, dtype=dtype) * 0.1 - gate_up_proj = torch.randn(experts, 2 * intermediate, hidden, device=device, dtype=dtype) * 0.02 - down_proj = torch.randn(experts, hidden, intermediate, device=device, dtype=dtype) * 0.02 - module = PackedExperts(gate_up_proj, down_proj).to(device) - - normalized_gate_up, normalized_down = _normalize_packed_expert_weights(module, dtype, hidden) - counts = torch.tensor(token_counts, device=device, dtype=torch.int64) - - with torch.no_grad(): - expected = reference_forward(inputs, token_counts, gate_up_proj, down_proj) - actual = gmm_forward(inputs, counts, normalized_gate_up, normalized_down) - # Reproduce the old DeepSeek-V4 bug: the square gate/up tensor was not transposed. - old_bug = gmm_forward(inputs, counts, gate_up_proj, down_proj.transpose(1, 2)) - torch.npu.synchronize() - - difference = (actual.float() - expected.float()).abs() - old_difference = (old_bug.float() - expected.float()).abs() - passed = torch.allclose(actual.float(), expected.float(), rtol=args.rtol, atol=args.atol) - report = { - 'device': str(device), - 'dtype': str(dtype), - 'input_shape': list(inputs.shape), - 'gate_up_shape_transformers': list(gate_up_proj.shape), - 'down_shape_transformers': list(down_proj.shape), - 'gate_up_shape_gmm': list(normalized_gate_up.shape), - 'down_shape_gmm': list(normalized_down.shape), - 'rtol': args.rtol, - 'atol': args.atol, - 'max_abs_diff': difference.max().item(), - 'mean_abs_diff': difference.mean().item(), - 'old_bug_max_abs_diff': old_difference.max().item(), - 'old_bug_mean_abs_diff': old_difference.mean().item(), - 'passed': bool(passed), - } - - args.output.parent.mkdir(parents=True, exist_ok=True) - args.output.write_text(json.dumps(report, indent=2), encoding='utf-8') - print(json.dumps(report, indent=2)) - print(f'Report saved to: {args.output.resolve()}') - if not passed: - raise SystemExit(1) - - -if __name__ == '__main__': - main() diff --git a/cookbook/client/server/transformer/read_me.md b/cookbook/client/server/transformer/read_me.md deleted file mode 100644 index cef5000fd..000000000 --- a/cookbook/client/server/transformer/read_me.md +++ /dev/null @@ -1,983 +0,0 @@ -# DeepSeek-V4 四层模型：FSDP/EP/NPU GMM 排查手册 - -本文档用于定位 Twinkle 加载 DeepSeek-V4 后基座生成异常的问题。整套检查只需要从正式 BF16 模型截取的前四层，不依赖模型能够正常回答自然语言。 - -排查目标是区分： - -1. checkpoint 或 Transformers 转换错误； -2. 两个节点加载出的源权重不一致； -3. FSDP node-local 权重分发错误； -4. EP 专家切分或 AllToAll 错误； -5. NPU grouped-matmul 专家权重方向错误。 - -四层模型无法形成有意义的自然语言回答是正常现象。所有结论必须依据权重诊断日志和最后一个位置的 logits，不能依据生成文本是否可读。 - -## 本次修改包含的文件 - -核心修复及诊断开关： - -- `src/twinkle/kernel/ops/moe/npu.py` - - 同时根据 `gate_up_proj` 和 `down_proj` 判断 Transformers `[E,out,in]` 与 GMM `[E,in,out]` 布局。 - - 修复 DeepSeek-V4 中 `hidden_size == 2 * moe_intermediate_size` 导致方阵 `gate_up_proj` 被错误识别的问题。 -- `src/twinkle/kernel/ops/ep/__init__.py` - - 新增 `TWINKLE_EP_FORCE_LOOP=1`，可强制使用逐专家 `F.linear` 参考实现。 -- `src/twinkle/model/transformers/moe/expert_parallel.py` - - 新增首轮 EP 路由、split、专家区间及输出诊断日志。 -- `src/twinkle/model/transformers/strategy/native_fsdp.py` - - 新增 node-local 源权重和 EP 本地专家切片诊断日志。 - -四层诊断配置及启动脚本： - -- `cookbook/client/server/transformer/server_config_dsv4_4layer_diag_no_ep.yaml` -- `cookbook/client/server/transformer/server_config_dsv4_4layer_diag_ep.yaml` -- `cookbook/client/server/transformer/run_dsv4_4layer_ep_diagnostic.sh` -- `cookbook/client/server/transformer/run_dsv4_0731_npu_2node_2npu.sh` - - 现在支持通过 `TWINKLE_SERVER_CONFIG_PATH` 选择配置。 - -诊断工具： - -- `cookbook/client/server/transformer/diagnostics/test_dsv4_npu_gmm_layout.py` -- `cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_ray.py` -- `cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_cli.py` -- `cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits.py` -- `cookbook/client/server/transformer/diagnostics/compare_dsv4_4layer_logits.py` - -单元测试： - -- `tests/kernel/ops/test_moe.py` - -## 1. 四层模型要求 - -必须使用从正式 DeepSeek-V4-Flash BF16 checkpoint 截取的前四层，不能使用随机初始化模型。配置至少需要保持： - -```text -num_hidden_layers=4 -hidden_size=4096 -moe_intermediate_size=2048 -n_routed_experts=256 -num_experts_per_tok=6 -dtype=bfloat16 -``` - -重点是保留： - -```text -hidden_size == 2 * moe_intermediate_size == 4096 -``` - -这正是原实现中方阵 `gate_up_proj` 布局误判的触发条件。 - -两个节点必须看到相同的绝对模型路径，并使用相同的 Twinkle、Transformers、PyTorch、torch-npu 和 CANN 版本。 - -启动前在两个节点分别检查： - -```bash -PROJECT_DIR=/opt/twinkle -MODEL_DIR=/nas/disk6/ljl/DeepSeek-V4-Flash-0731-BF16-4layers - -cd "$PROJECT_DIR" -git rev-parse HEAD -test -f "$MODEL_DIR/config.json" -test -f "$MODEL_DIR/model.safetensors.index.json" -python3 - <<'PY' -import torch -import transformers -import torch_npu - -print('torch=', torch.__version__) -print('torch_npu=', torch_npu.__version__) -print('transformers=', transformers.__version__) -PY -``` - -两个节点的输出必须一致。 - -## 2. 避免占用 Pod ephemeral-storage - -Ray 的 Unix socket 路径必须短，但缓存又不应写入只有 10 GiB 的容器临时盘。已知容器的 `/dev/shm` 有 800 GiB 时，建议把 Ray 临时目录放到共享内存；日志和结果仍写入存储卷： - -```bash -mkdir -p /dev/shm/rh -mkdir -p /nas/disk6/ljl/dsv4_ep_diag/tmp -mkdir -p /nas/disk6/ljl/dsv4_ep_diag/logs -mkdir -p /nas/disk6/ljl/dsv4_ep_diag/results -``` - -之后使用： - -```bash -export RAY_TMPDIR=/dev/shm/rh -export TMPDIR=/nas/disk6/ljl/dsv4_ep_diag/tmp -export RAY_ROTATION_MAX_BYTES=20971520 -export RAY_ROTATION_BACKUP_COUNT=1 -``` - -不要把 Ray 临时目录设成很长的路径，否则可能再次触发 AF_UNIX 107 字节路径上限。 - -## 3. 先做单卡 NPU GMM 数值测试 - -该步骤不启动 Ray，也不加载四层模型。它使用保持 DeepSeek-V4 比例的微型专家，直接比较： - -```text -F.linear 参考结果 -vs -npu_grouped_matmul 结果 -``` - -在一台 NPU 机器执行： - -```bash -cd /opt/twinkle - -ASCEND_RT_VISIBLE_DEVICES=0 \ -PYTHONPATH=/opt/twinkle/src \ -python3 cookbook/client/server/transformer/diagnostics/test_dsv4_npu_gmm_layout.py \ - --device npu:0 \ - --output /nas/disk6/ljl/dsv4_ep_diag/results/npu_gmm_layout.json -``` - -成功时： - -```text -passed=true -max_abs_diff 和 mean_abs_diff 在 BF16 允许范围内 -old_bug_max_abs_diff 明显大于修复后的 max_abs_diff -``` - -如果这里失败，不要启动分布式服务，先处理 NPU GMM 与 `F.linear` 的数值差异。 - -## 3A. 多机 Ray CLI 基准（推荐） - -这里的“CLI”表示在 Head 节点运行一个普通 Python driver，而不是必须使用 -torchrun。driver 直接通过 Ray 创建两节点 NPU 模型 Actor： - -```text -Python driver - -> Ray placement group - -> 两节点、每节点两个 MultiLoraTransformersModel Actor - -> native_fsdp + HCCL + EP - -> logits 返回 driver -``` - -该模式不启动 Ray Serve、GatewayServer 或 ModelManagement，不创建 HTTP session -或租户，因此不受服务端槽位、心跳和接口限流影响。模型内部会创建一个禁用 delta -的诊断 adapter，以复用生产 `MultiLoraTransformersModel` 在 FSDP 包装前安装槽位和 -对齐 dtype 的相同初始化路径；该 adapter 不登记到 ServerState 容量中。该模式仍然 -覆盖 Ray 多机 Actor、node-local 权重加载、`_broadcast_sharded_state_dict`、FSDP -和 EP。 - -先创建短临时目录和持久化结果目录： - -```bash -mkdir -p /dev/shm/rh -mkdir -p /nas/disk6/ljl/dsv4_ep_diag/tmp -mkdir -p /nas/disk6/ljl/dsv4_ep_diag/logs -mkdir -p /nas/disk6/ljl/dsv4_ep_diag/results -``` - -Head 节点： - -```bash -cd /opt/twinkle - -export HEAD_IP=172.61.10.254 -export ASCEND_RT_VISIBLE_DEVICES=0,1 -export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1 -export HCCL_SOCKET_IFNAME=eth0 -export GLOO_SOCKET_IFNAME=eth0 -export PYTHONPATH=/opt/twinkle/src -export TMPDIR=/nas/disk6/ljl/dsv4_ep_diag/tmp - -ray stop --force || true -ray start \ - --head \ - --node-ip-address="$HEAD_IP" \ - --port=6379 \ - --num-cpus="$(nproc)" \ - --resources='{"NPU": 2}' \ - --temp-dir=/dev/shm/rh \ - --disable-usage-stats \ - --include-dashboard=false -``` - -Worker 节点： - -```bash -cd /opt/twinkle - -export HEAD_IP=172.61.10.254 -export NODE_IP=172.61.12.251 -export ASCEND_RT_VISIBLE_DEVICES=0,1 -export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1 -export HCCL_SOCKET_IFNAME=eth0 -export GLOO_SOCKET_IFNAME=eth0 -export PYTHONPATH=/opt/twinkle/src -export TMPDIR=/nas/disk6/ljl/dsv4_ep_diag/tmp - -ray stop --force || true -ray start \ - --address="$HEAD_IP:6379" \ - --node-ip-address="$NODE_IP" \ - --num-cpus="$(nproc)" \ - --resources='{"NPU": 2}' \ - --temp-dir=/dev/shm/rh \ - --disable-usage-stats -``` - -Worker 加入后，只在 Head 节点运行 driver： - -```bash -cd /opt/twinkle - -export DSV4_MODEL_ID=/nas/disk6/ljl/DeepSeek-V4-Flash-0731-BF16-4layers -export ASCEND_RT_VISIBLE_DEVICES=0,1 -export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1 -export HCCL_SOCKET_IFNAME=eth0 -export GLOO_SOCKET_IFNAME=eth0 -export PYTHONPATH=/opt/twinkle/src -export TMPDIR=/nas/disk6/ljl/dsv4_ep_diag/tmp - -python3 cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_ray.py \ - --ray-address auto \ - --world-size 4 \ - --nproc-per-node 2 \ - --mode no_ep \ - --output-dir /nas/disk6/ljl/dsv4_ep_diag/results \ - 2>&1 | tee /nas/disk6/ljl/dsv4_ep_diag/logs/ray_no_ep.log -``` - -完成后，Ray 集群保持运行。分别重新执行 driver，只修改模式和日志名： - -```bash -python3 cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_ray.py \ - --ray-address auto --world-size 4 --nproc-per-node 2 \ - --mode ep_loop \ - --output-dir /nas/disk6/ljl/dsv4_ep_diag/results \ - 2>&1 | tee /nas/disk6/ljl/dsv4_ep_diag/logs/ray_ep_loop.log - -python3 cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_ray.py \ - --ray-address auto --world-size 4 --nproc-per-node 2 \ - --mode ep_gmm \ - --output-dir /nas/disk6/ljl/dsv4_ep_diag/results \ - 2>&1 | tee /nas/disk6/ljl/dsv4_ep_diag/logs/ray_ep_gmm.log -``` - -每个 driver 结束时只删除自己创建的模型 Actor 和 placement group，不执行 -`ray stop`。输出文件为： - -```text -ray_no_ep_last_logits.pt/json -ray_ep_loop_last_logits.pt/json -ray_ep_gmm_last_logits.pt/json -``` - -如果首次 forward 报错： - -```text -AssertionError: FSDP expects uniform original parameter dtype but got -{torch.bfloat16, torch.float32} -``` - -说明运行的是旧版 Ray probe：它实例化了纯 `TransformersModel`，没有进入生产 -Multi-LoRA 的槽位安装和 `_ensure_lora_dtype()` 路径。使用更新后的 -`probe_dsv4_4layer_logits_ray.py`；它直接构造 `MultiLoraTransformersModel`，先安装 -禁用的诊断 adapter，再进入 `_lazy_wrap_model()`。无需修改公共 -`TransformersModel._lazy_wrap_model()`，Ray 集群也无需重启。 - -比较方式： - -```bash -python3 cookbook/client/server/transformer/diagnostics/compare_dsv4_4layer_logits.py \ - /nas/disk6/ljl/dsv4_ep_diag/results/ray_no_ep_last_logits.pt \ - /nas/disk6/ljl/dsv4_ep_diag/results/ray_ep_loop_last_logits.pt \ - --output /nas/disk6/ljl/dsv4_ep_diag/results/compare_ray_no_ep_vs_ep_loop.json - -python3 cookbook/client/server/transformer/diagnostics/compare_dsv4_4layer_logits.py \ - /nas/disk6/ljl/dsv4_ep_diag/results/ray_ep_loop_last_logits.pt \ - /nas/disk6/ljl/dsv4_ep_diag/results/ray_ep_gmm_last_logits.pt \ - --output /nas/disk6/ljl/dsv4_ep_diag/results/compare_ray_ep_loop_vs_ep_gmm.json -``` - -## 3B. 可选的本地/torchrun 基准 - -CLI 模式直接在 torchrun 进程中构造 `MultiLoraTransformersModel`，不启动 Ray、 -Gateway 或 ModelManagement，也不会创建 session、租户或登记服务端 LoRA 容量。 -它会在每个进程内安装一个禁用 delta 的诊断 adapter，使 dtype 对齐和 FSDP 包装 -顺序与生产服务一致。该步骤验证： - -```text -Transformers checkpoint 转换 -native_fsdp 包装和权重加载 -EP 专家切分与通信 -loop/GMM 专家计算 -``` - -它不会验证 Ray Actor 和服务端 `_broadcast_sharded_state_dict` 路径，因此 CLI -通过后仍需执行后面的 C/S 对照。 - -### 单节点四张 NPU - -```bash -cd /opt/twinkle - -export DSV4_MODEL_ID=/nas/disk6/ljl/DeepSeek-V4-Flash-0731-BF16-4layers -export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3 -export HCCL_SOCKET_IFNAME=bond0 -export GLOO_SOCKET_IFNAME=bond0 -export PYTHONPATH=/opt/twinkle/src - -torchrun --standalone --nproc-per-node=4 \ - cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_cli.py \ - --mode no_ep \ - --output-dir /nas/disk6/ljl/dsv4_ep_diag/results - -torchrun --standalone --nproc-per-node=4 \ - cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_cli.py \ - --mode ep_loop \ - --output-dir /nas/disk6/ljl/dsv4_ep_diag/results - -torchrun --standalone --nproc-per-node=4 \ - cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_cli.py \ - --mode ep_gmm \ - --output-dir /nas/disk6/ljl/dsv4_ep_diag/results -``` - -每种模式必须启动一组新的 torchrun 进程，不能在同一 Python 进程内切换 -`TWINKLE_EP_FORCE_LOOP`。 - -### 两节点、每节点两张 NPU - -以下示例中 Head 为 `172.61.10.254`，Worker 为 `172.61.12.251`。先在 Head -执行 rank 0 命令，它会等待 Worker；然后在 Worker 执行 rank 1 命令。 - -Head： - -```bash -cd /opt/twinkle - -export DSV4_MODEL_ID=/nas/disk6/ljl/DeepSeek-V4-Flash-0731-BF16-4layers -export ASCEND_RT_VISIBLE_DEVICES=0,1 -export HCCL_SOCKET_IFNAME=eth0 -export GLOO_SOCKET_IFNAME=eth0 -export PYTHONPATH=/opt/twinkle/src - -torchrun \ - --nnodes=2 \ - --nproc-per-node=2 \ - --node-rank=0 \ - --master-addr=172.61.10.254 \ - --master-port=29610 \ - cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_cli.py \ - --mode no_ep \ - --output-dir /nas/disk6/ljl/dsv4_ep_diag/results -``` - -Worker： - -```bash -cd /opt/twinkle - -export DSV4_MODEL_ID=/nas/disk6/ljl/DeepSeek-V4-Flash-0731-BF16-4layers -export ASCEND_RT_VISIBLE_DEVICES=0,1 -export HCCL_SOCKET_IFNAME=eth0 -export GLOO_SOCKET_IFNAME=eth0 -export PYTHONPATH=/opt/twinkle/src - -torchrun \ - --nnodes=2 \ - --nproc-per-node=2 \ - --node-rank=1 \ - --master-addr=172.61.10.254 \ - --master-port=29610 \ - cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_cli.py \ - --mode no_ep \ - --output-dir /nas/disk6/ljl/dsv4_ep_diag/results -``` - -完成 `no_ep` 后，把两个节点命令中的 `--mode no_ep` 同时改成 -`--mode ep_loop` 再运行一次，然后改成 `--mode ep_gmm`。上一组 torchrun 已经 -完全退出后可以复用端口 `29610`；如果仍有残留进程，则先终止残留进程或换一个 -未占用端口。 - -只使用一张 NPU 时只能验证 `no_ep`： - -```bash -ASCEND_RT_VISIBLE_DEVICES=0 \ -PYTHONPATH=/opt/twinkle/src \ -python3 cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_cli.py \ - --mode no_ep \ - --model-id /nas/disk6/ljl/DeepSeek-V4-Flash-0731-BF16-4layers \ - --output-dir /nas/disk6/ljl/dsv4_ep_diag/results -``` - -CLI 输出： - -```text -cli_no_ep_last_logits.pt/json -cli_ep_loop_last_logits.pt/json -cli_ep_gmm_last_logits.pt/json -``` - -报告中的 `max_rank_diff` 应接近 0。使用现有比较脚本比较 CLI 三组结果： - -```bash -python3 cookbook/client/server/transformer/diagnostics/compare_dsv4_4layer_logits.py \ - /nas/disk6/ljl/dsv4_ep_diag/results/cli_no_ep_last_logits.pt \ - /nas/disk6/ljl/dsv4_ep_diag/results/cli_ep_loop_last_logits.pt \ - --output /nas/disk6/ljl/dsv4_ep_diag/results/compare_cli_no_ep_vs_ep_loop.json - -python3 cookbook/client/server/transformer/diagnostics/compare_dsv4_4layer_logits.py \ - /nas/disk6/ljl/dsv4_ep_diag/results/cli_ep_loop_last_logits.pt \ - /nas/disk6/ljl/dsv4_ep_diag/results/cli_ep_gmm_last_logits.pt \ - --output /nas/disk6/ljl/dsv4_ep_diag/results/compare_cli_ep_loop_vs_ep_gmm.json -``` - -### 比较 memory_efficient_init 开关后的实际权重 - -最终 logits 不同只能证明初始化路径影响了结果。要定位具体参数，使用相同的 -`ep_loop` 计算路径分别开启和关闭 `memory_efficient_init`，并增加 -`--dump-weight-fingerprints`。诊断只读取每个 rank 的本地参数，并从每个参数均匀 -采样最多 4096 个值，不会将完整权重保存到磁盘，也不会为了展平非连续参数而复制 -完整权重。 - -开启： - -```bash -torchrun --standalone --nproc-per-node=4 \ - cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_cli.py \ - --mode ep_loop \ - --dump-weight-fingerprints \ - --output-dir /nas/disk6/ljl/dsv4_ep_diag/results/memory_efficient_on -``` - -关闭： - -```bash -torchrun --standalone --nproc-per-node=4 \ - cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_cli.py \ - --mode ep_loop \ - --disable-memory-efficient-init \ - --dump-weight-fingerprints \ - --output-dir /nas/disk6/ljl/dsv4_ep_diag/results/memory_efficient_off -``` - -每次运行会按 rank 生成： - -```text -cli_ep_loop_weight_fingerprints_rank0.json -cli_ep_loop_weight_fingerprints_rank1.json -cli_ep_loop_weight_fingerprints_rank2.json -cli_ep_loop_weight_fingerprints_rank3.json -``` - -以关闭状态作为 reference 进行比较： - -```bash -python3 cookbook/client/server/transformer/diagnostics/compare_dsv4_weight_fingerprints.py \ - /nas/disk6/ljl/dsv4_ep_diag/results/memory_efficient_off \ - /nas/disk6/ljl/dsv4_ep_diag/results/memory_efficient_on \ - --mode ep_loop \ - --output /nas/disk6/ljl/dsv4_ep_diag/results/compare_memory_init_weights.json \ - || true -``` - -重点查看比较报告中的以下字段： - -```text -metadata_mismatches shape、stride、dtype 或连续性不同 -value_mismatches 参数采样值不同 -first_difference 第一个不同采样值的本地 flat index 和多维坐标 -rank 不同权重所在的 EP/FSDP rank -name 不同权重的完整参数名 -``` - -如果第一个不一致参数是 `embed_tokens`、attention、norm 或 `lm_head`，优先检查 -普通 FSDP DTensor 重建。如果只有 `mlp.experts.gate_up_proj/down_proj` 不一致，优先 -检查 pre-EP state capture 和 `_scatter_ep_expert_tensor()` 的专家切分。 - -还可以将 `cli__last_logits.pt` 与 C/S 产生的 -`_last_logits.pt` 比较。CLI 一致而 C/S 不一致时，重点排查 Ray/服务端权重 -分发；CLI 本身已经不一致时，重点排查 FSDP/EP/GMM。 - -## 4. 三组服务端对照模式 - -三组模式只改变 EP 和专家计算方式： - -| 模式 | FSDP | EP | 专家计算 | -|---|---:|---:|---| -| `no_ep` | 开启 | 关闭 | Transformers 原始前向 | -| `ep_loop` | 开启 | 开启 | 强制逐专家 `F.linear` | -| `ep_gmm` | 开启 | 开启 | NPU grouped-matmul | - -两份 YAML 除 EP 设置外保持一致： - -```text -两节点 -每节点 2 张 NPU -world_size=4 -fsdp_size=4 -ep_size=4（仅 EP 模式） -memory_efficient_init=true -max_loras=1 -max_length=512 -``` - -EP=4 时，256 个专家的预期区间为： - -```text -rank 0 / ep_rank 0: experts 0..63 -rank 1 / ep_rank 1: experts 64..127 -rank 2 / ep_rank 2: experts 128..191 -rank 3 / ep_rank 3: experts 192..255 -``` - -## 5. 两个节点的公共环境 - -以下示例使用： - -```text -Head: 172.61.10.254 -Worker: 172.61.12.251 -网卡: eth0 -``` - -如果实际环境使用 `bond0` 或其他地址，只修改环境变量，不需要修改脚本或 YAML。 - -Head 节点执行： - -```bash -cd /opt/twinkle - -export DSV4_MODEL_ID=/nas/disk6/ljl/DeepSeek-V4-Flash-0731-BF16-4layers -export HEAD_IP=172.61.10.254 -export WORKER_IP=172.61.12.251 -export NODE_IP=172.61.10.254 -export NETWORK_IFACE=eth0 -export ASCEND_RT_VISIBLE_DEVICES=0,1 -export RESET_RAY=1 -export RAY_TMPDIR=/dev/shm/rh -export TMPDIR=/nas/disk6/ljl/dsv4_ep_diag/tmp -export RAY_ROTATION_MAX_BYTES=20971520 -export RAY_ROTATION_BACKUP_COUNT=1 -``` - -Worker 节点执行： - -```bash -cd /opt/twinkle - -export DSV4_MODEL_ID=/nas/disk6/ljl/DeepSeek-V4-Flash-0731-BF16-4layers -export HEAD_IP=172.61.10.254 -export WORKER_IP=172.61.12.251 -export NODE_IP=172.61.12.251 -export NETWORK_IFACE=eth0 -export ASCEND_RT_VISIBLE_DEVICES=0,1 -export RESET_RAY=1 -export RAY_TMPDIR=/dev/shm/rh -export TMPDIR=/nas/disk6/ljl/dsv4_ep_diag/tmp -export RAY_ROTATION_MAX_BYTES=20971520 -export RAY_ROTATION_BACKUP_COUNT=1 -``` - -必须确保两个节点选择相同模式。 - -## 6. 模式一：启动 no-EP 基准 - -先在 Head 节点启动；脚本会等待 Worker 加入： - -```bash -nohup bash cookbook/client/server/transformer/run_dsv4_4layer_ep_diagnostic.sh head no_ep \ - >/nas/disk6/ljl/dsv4_ep_diag/logs/no_ep_head.log 2>&1 /nas/disk6/ljl/dsv4_ep_diag/logs/no_ep_worker.log 2>&1 /nas/disk6/ljl/dsv4_ep_diag/logs/ep_loop_head.log 2>&1 /nas/disk6/ljl/dsv4_ep_diag/logs/ep_loop_worker.log 2>&1 /nas/disk6/ljl/dsv4_ep_diag/logs/ep_gmm_head.log 2>&1 /nas/disk6/ljl/dsv4_ep_diag/logs/ep_gmm_worker.log 2>&1 /nas/disk6/ljl/dsv4_ep_diag/results/ep_diagnostics.txt -``` - -重点检查以下日志。 - -### 11.1 两个节点的源权重 - -每个 node-local source rank 都会输出： - -```text -local_source= -param=<参数名> -full_shape=<完整专家形状> -source_preview=<采样值> -``` - -相同参数在两个 node-local source 上必须满足： - -```text -full_shape 相同 -source_preview 相同 -``` - -如果不同，优先检查两个节点： - -- 是否使用相同模型目录； -- 模型文件是否一致； -- Transformers/torch版本是否一致； -- 是否实际运行同一个 Twinkle commit。 - -### 11.2 EP 专家范围 - -每个rank会输出： - -```text -rank= -ep_rank= -expert_range=[start,end) -local_shape=<本地权重形状> -local_preview=<本地切片采样值> -``` - -EP=4时必须严格对应： - -```text -rank0 -> [0,64) -rank1 -> [64,128) -rank2 -> [128,192) -rank3 -> [192,256) -``` - -本地权重预期形状： - -```text -gate_up_proj: [64,4096,4096] -down_proj: [64,4096,2048] -``` - -### 11.3 Router与AllToAll split - -每层第一次前向会输出： - -```text -selected_experts -routing_weights -input_splits -output_splits -output_finite -``` - -固定输入下，各rank在AllToAll前的 `selected_experts` 和 `routing_weights` 应一致。代码还会验证： - -```text -sum(input_splits) == token_count * num_experts_per_tok -``` - -如果该条件不成立，诊断模式会直接报错退出。 - -## 12. 单元测试和静态检查 - -在安装了项目依赖的环境执行： - -```bash -cd /opt/twinkle - -python3 -m pytest -q tests/kernel/ops/test_moe.py - -python3 -m py_compile \ - cookbook/client/server/transformer/diagnostics/test_dsv4_npu_gmm_layout.py \ - cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_ray.py \ - cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits_cli.py \ - cookbook/client/server/transformer/diagnostics/probe_dsv4_4layer_logits.py \ - cookbook/client/server/transformer/diagnostics/compare_dsv4_4layer_logits.py \ - cookbook/client/server/transformer/diagnostics/compare_dsv4_weight_fingerprints.py - -bash -n \ - cookbook/client/server/transformer/run_dsv4_0731_npu_2node_2npu.sh \ - cookbook/client/server/transformer/run_dsv4_4layer_ep_diagnostic.sh -``` - -## 13. 完整模型的最终验证 - -只有满足以下条件后才启动43层完整模型： - -```text -单卡 NPU GMM 与 F.linear 一致 -no_ep logits ≈ ep_loop logits -ep_loop logits ≈ ep_gmm logits -两个 node-local source 的权重采样一致 -EP 专家范围正确 -路由 split 不变量通过 -``` - -完整模型启动时应关闭诊断开关： - -```bash -export TWINKLE_EP_FORCE_LOOP=0 -export TWINKLE_EP_DIAGNOSTICS=0 -``` - -然后按顺序验证： - -1. 不训练、不注册旧LoRA，先检查基座生成； -2. 注册全零LoRA，确认它与 `disable_lora=True` 逐token一致； -3. 修复前在错误EP前向下训练的LoRA不能作为正确性基准； -4. 基座确认正常后重新训练LoRA； -5. 对比重新训练后的 base/logprob/生成结果。 - -## 14. 清理 - -诊断结束后在两个节点执行： - -```bash -ray stop --force -``` - -结果位于共享卷： - -```text -/nas/disk6/ljl/dsv4_ep_diag/results -/nas/disk6/ljl/dsv4_ep_diag/logs -``` - -确认不再需要后可手动清理共享卷中的诊断日志和结果。`/dev/shm/rh` 仅用于当前 Pod 的 Ray 临时文件，Pod 删除后不会保留。 - -## 15. 路由 Management 类型 NameError - -如果应用注册路由时出现以下任意错误： - -```text -NameError: name 'GatewayServer' is not defined -NameError: name 'ModelManagement' is not defined -NameError: name 'SamplerManagement' is not defined -NameError: name 'ProcessorManagement' is not defined -``` - -说明运行版本中的 handler 只在 `TYPE_CHECKING` 分支导入了对应的 Management -类型，但 FastAPI 在运行时解析了路由参数的类型注解。这发生在服务配置加载阶段， -与模型权重、FSDP 和 EP 无关。 - -当前源码已在 Gateway、Model、Sampler 和 Processor 的 handler 中提供无循环导入 -的运行时类型别名，包括： - -```text -src/twinkle/server/gateway/openai_handlers.py -src/twinkle/server/gateway/tinker_handlers.py -src/twinkle/server/gateway/twinkle_handlers.py -src/twinkle/server/model/tinker_handlers.py -src/twinkle/server/model/twinkle_handlers.py -src/twinkle/server/sampler/tinker_handlers.py -src/twinkle/server/sampler/twinkle_handlers.py -src/twinkle/server/processor/twinkle_handlers.py -``` - -将修复后的源码同步到 head 和 worker 使用的相同 Twinkle 目录，然后重新启动服务。 -可先做静态检查： - -```bash -cd /opt/twinkle -python3 -m compileall -q \ - src/twinkle/server/gateway \ - src/twinkle/server/model \ - src/twinkle/server/sampler \ - src/twinkle/server/processor -``` - -如果错误仍然存在，先确认容器实际导入的源码位置： - -```bash -cd /opt/twinkle -PYTHONPATH=/opt/twinkle/src python3 - <<'PY' -import inspect -import twinkle.server.gateway.openai_handlers as module - -print(inspect.getfile(module)) -print('GatewayServer runtime alias:', module.GatewayServer) -PY -``` - -输出文件应来自 `/opt/twinkle/src/twinkle/server/gateway/`，且不应再次出现 -`NameError`。 diff --git a/cookbook/client/server/transformer/run_dsv4_0731_npu.sh b/cookbook/client/server/transformer/run_dsv4_0731_npu.sh deleted file mode 100755 index dd2eea223..000000000 --- a/cookbook/client/server/transformer/run_dsv4_0731_npu.sh +++ /dev/null @@ -1,67 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail - -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" -PROJECT_DIR="$(cd -- "$SCRIPT_DIR/../../../.." && pwd)" -cd "$PROJECT_DIR" - -if [[ -f /usr/local/Ascend/ascend-toolkit/set_env.sh ]]; then - # shellcheck disable=SC1091 - source /usr/local/Ascend/ascend-toolkit/set_env.sh -fi - -export PYTHONPATH="$PROJECT_DIR/src${PYTHONPATH:+:$PYTHONPATH}" -export ASCEND_RT_VISIBLE_DEVICES="${ASCEND_RT_VISIBLE_DEVICES:-0,1}" -export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1 -export TWINKLE_TRUST_REMOTE_CODE=1 -export TWINKLE_FAIL_FAST=1 -export TOKENIZERS_PARALLELISM=true -export HCCL_CONNECT_TIMEOUT="${HCCL_CONNECT_TIMEOUT:-7200}" -export HCCL_EXEC_TIMEOUT="${HCCL_EXEC_TIMEOUT:-0}" - -MODEL_PATH=/nas/disk1/random-deepseek-v4-4b -CONFIG_PATH="$SCRIPT_DIR/server_config_dsv4_0731_npu.yaml" - -test -f "$MODEL_PATH/config.json" -test -f "$MODEL_PATH/tokenizer.json" - -python3 - <<'PY' -import os -import torch -import torch_npu # noqa: F401 - -visible = [item.strip() for item in os.environ['ASCEND_RT_VISIBLE_DEVICES'].split(',') if item.strip()] -if len(visible) != 2: - raise SystemExit(f'需要两张 NPU，当前 ASCEND_RT_VISIBLE_DEVICES={visible}') -if not torch.npu.is_available(): - raise SystemExit('torch.npu.is_available() 为 False') -if torch.npu.device_count() < 2: - raise SystemExit(f'当前进程只能看到 {torch.npu.device_count()} 张 NPU') -print(f'NPU 检查通过：visible={visible}, device_count={torch.npu.device_count()}') -PY - -if ! ray status >/dev/null 2>&1; then - ray start \ - --head \ - --num-cpus="${TWINKLE_RAY_CPUS:-8}" \ - --resources='{"NPU": 2}' \ - --disable-usage-stats \ - --include-dashboard=false -fi - -python3 - <<'PY' -import ray - -ray.init(address='auto', logging_level='ERROR') -npu_count = float(ray.cluster_resources().get('NPU', 0)) -ray.shutdown() -if npu_count < 2: - raise SystemExit( - f'当前 Ray 集群只有 {npu_count:g} 个 NPU 资源。' - '请确认没有其他任务占用后执行 ray stop --force，再重新运行本脚本。' - ) -print(f'Ray NPU 资源检查通过：NPU={npu_count:g}') -PY - -python3 -m twinkle.server check-config --config "$CONFIG_PATH" -exec python3 "$SCRIPT_DIR/server_dsv4_0731_npu.py" diff --git a/cookbook/client/server/transformer/run_dsv4_0731_npu_2node_2npu.sh b/cookbook/client/server/transformer/run_dsv4_0731_npu_2node_2npu.sh deleted file mode 100755 index f7fadf5ca..000000000 --- a/cookbook/client/server/transformer/run_dsv4_0731_npu_2node_2npu.sh +++ /dev/null @@ -1,170 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail - -# Head node (172.61.10.254): -# bash cookbook/client/server/transformer/run_dsv4_0731_npu_2node_2npu.sh head -# -# Worker node (172.61.12.251): -# bash cookbook/client/server/transformer/run_dsv4_0731_npu_2node_2npu.sh worker -# -# Override NETWORK_IFACE when the Ray/HCCL interface is not eth0. Set -# RESET_RAY=1 when intentionally replacing an existing local Ray runtime. - -ROLE="${1:-}" -if [[ "$ROLE" != "head" && "$ROLE" != "worker" ]]; then - echo "Usage: $0 {head|worker}" >&2 - exit 2 -fi - -DEFAULT_HEAD_IP="172.61.10.254" -DEFAULT_WORKER_IP="172.61.12.251" -HEAD_IP="${HEAD_IP:-$DEFAULT_HEAD_IP}" -WORKER_IP="${WORKER_IP:-$DEFAULT_WORKER_IP}" - -if [[ "$ROLE" == "head" ]]; then - NODE_IP="${NODE_IP:-$HEAD_IP}" - NODE_RANK=0 -else - NODE_IP="${NODE_IP:-$WORKER_IP}" - NODE_RANK=1 -fi - -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" -PROJECT_DIR="$(cd -- "$SCRIPT_DIR/../../../.." && pwd)" -CONFIG_PATH="${TWINKLE_SERVER_CONFIG_PATH:-$SCRIPT_DIR/server_config_dsv4_0731_npu_2node_2npu.yaml}" -if [[ "$CONFIG_PATH" != /* ]]; then - CONFIG_PATH="$PROJECT_DIR/$CONFIG_PATH" -fi - -export DSV4_MODEL_ID="${DSV4_MODEL_ID:-hf://deepseek-ai/DeepSeek-V4-Flash-0731}" -DATASET_PATH="${DATASET_ID:-/model/ljl/dataset/self-cognition.jsonl}" - -NPU_PER_NODE=2 -NNODES=2 -TOTAL_NPUS=$((NPU_PER_NODE * NNODES)) -RAY_PORT="${RAY_PORT:-6379}" -RAY_CPUS_PER_NODE="${TWINKLE_RAY_CPUS:-$(nproc)}" -CLUSTER_WAIT_SECONDS="${CLUSTER_WAIT_SECONDS:-1800}" -NETWORK_IFACE="${NETWORK_IFACE:-eth0}" - -if [[ -f /usr/local/Ascend/ascend-toolkit/set_env.sh ]]; then - # shellcheck disable=SC1091 - source /usr/local/Ascend/ascend-toolkit/set_env.sh -fi - -cd "$PROJECT_DIR" -export PYTHONPATH="$PROJECT_DIR/src${PYTHONPATH:+:$PYTHONPATH}" -export ASCEND_RT_VISIBLE_DEVICES="${ASCEND_RT_VISIBLE_DEVICES:-0,1}" -export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1 -export TWINKLE_TRUST_REMOTE_CODE=1 -export TWINKLE_FAIL_FAST=1 -export TWINKLE_EP_FORCE_LOOP="${TWINKLE_EP_FORCE_LOOP:-0}" -export TWINKLE_EP_DIAGNOSTICS="${TWINKLE_EP_DIAGNOSTICS:-0}" -export TOKENIZERS_PARALLELISM=true -export GLOO_SOCKET_IFNAME="$NETWORK_IFACE" -export HCCL_SOCKET_IFNAME="$NETWORK_IFACE" -export HCCL_CONNECT_TIMEOUT="${HCCL_CONNECT_TIMEOUT:-7200}" -export HCCL_EXEC_TIMEOUT="${HCCL_EXEC_TIMEOUT:-0}" -export HCCL_IF_BASE_PORT="${HCCL_IF_BASE_PORT:-20000}" -export NODE_RANK - -IFS=',' read -r -a VISIBLE_NPUS <<< "$ASCEND_RT_VISIBLE_DEVICES" -if [[ "${#VISIBLE_NPUS[@]}" -ne "$NPU_PER_NODE" ]]; then - echo "Expected $NPU_PER_NODE visible NPUs, got ASCEND_RT_VISIBLE_DEVICES=$ASCEND_RT_VISIBLE_DEVICES" >&2 - exit 1 -fi - -test -f "$CONFIG_PATH" -if [[ "${SKIP_DATASET_CHECK:-0}" != "1" ]]; then - test -f "$DATASET_PATH" -fi - -if [[ "$DSV4_MODEL_ID" == hf://* || "$DSV4_MODEL_ID" == ms://* ]]; then - echo "Model will be downloaded through the configured Hub backend: $DSV4_MODEL_ID" -elif [[ -d "$DSV4_MODEL_ID" ]]; then - test -f "$DSV4_MODEL_ID/config.json" - test -f "$DSV4_MODEL_ID/tokenizer.json" -else - echo "Invalid DSV4_MODEL_ID: use an existing local directory or an explicit hf:// or ms:// ID." >&2 - exit 1 -fi - -python3 - "$NPU_PER_NODE" <<'PY' -import sys -import torch -import torch_npu # noqa: F401 - -expected = int(sys.argv[1]) -if not torch.npu.is_available(): - raise SystemExit('torch.npu.is_available() is False') -actual = torch.npu.device_count() -if actual < expected: - raise SystemExit(f'Expected at least {expected} visible NPUs, got {actual}') -print(f'Ascend check passed: visible NPU count={actual}') -PY - -if [[ "${RESET_RAY:-0}" == "1" ]]; then - ray stop --force || true -fi - -if ray status >/dev/null 2>&1; then - echo "Ray is already running on this node; set RESET_RAY=1 to replace it." -elif [[ "$ROLE" == "head" ]]; then - ray start \ - --head \ - --node-ip-address="$HEAD_IP" \ - --port="$RAY_PORT" \ - --num-cpus="$RAY_CPUS_PER_NODE" \ - --resources="{\"NPU\": $NPU_PER_NODE}" \ - --disable-usage-stats \ - --include-dashboard=false -else - ray start \ - --address="$HEAD_IP:$RAY_PORT" \ - --node-ip-address="$NODE_IP" \ - --num-cpus="$RAY_CPUS_PER_NODE" \ - --resources="{\"NPU\": $NPU_PER_NODE}" \ - --disable-usage-stats -fi - -if [[ "$ROLE" == "worker" ]]; then - echo "Worker $NODE_IP joined Ray at $HEAD_IP:$RAY_PORT with $NPU_PER_NODE NPU resources." - exit 0 -fi - -python3 - "$TOTAL_NPUS" "$NNODES" "$NPU_PER_NODE" "$CLUSTER_WAIT_SECONDS" <<'PY' -import sys -import time -import ray - -expected_total = int(sys.argv[1]) -expected_nodes = int(sys.argv[2]) -expected_per_node = int(sys.argv[3]) -timeout = int(sys.argv[4]) -deadline = time.monotonic() + timeout - -ray.init(address='auto', logging_level='ERROR') -try: - while True: - alive_nodes = [node for node in ray.nodes() if node.get('Alive', True)] - npu_nodes = [ - node for node in alive_nodes - if float(node.get('Resources', {}).get('NPU', 0)) >= expected_per_node - ] - total = int(ray.cluster_resources().get('NPU', 0)) - print(f'Waiting for cluster: NPU={total}/{expected_total}, NPU nodes={len(npu_nodes)}/{expected_nodes}') - if total >= expected_total and len(npu_nodes) >= expected_nodes: - break - if time.monotonic() >= deadline: - raise SystemExit( - f'Timed out waiting for {expected_nodes} nodes and {expected_total} total NPU resources') - time.sleep(5) -finally: - ray.shutdown() -PY - -python3 -m twinkle.server check-config --config "$CONFIG_PATH" -echo "Twinkle config: $CONFIG_PATH" -echo "EP diagnostics: force_loop=$TWINKLE_EP_FORCE_LOOP diagnostics=$TWINKLE_EP_DIAGNOSTICS" -echo "Launching Twinkle Server at http://$HEAD_IP:8000" -exec python3 -m twinkle.server launch --config "$CONFIG_PATH" diff --git a/cookbook/client/server/transformer/run_dsv4_0731_npu_multinode.sh b/cookbook/client/server/transformer/run_dsv4_0731_npu_multinode.sh index b2547993c..37123a9fa 100755 --- a/cookbook/client/server/transformer/run_dsv4_0731_npu_multinode.sh +++ b/cookbook/client/server/transformer/run_dsv4_0731_npu_multinode.sh @@ -26,7 +26,6 @@ SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" PROJECT_DIR="$(cd -- "$SCRIPT_DIR/../../../.." && pwd)" CONFIG_PATH="$SCRIPT_DIR/server_config_dsv4_0731_npu_multinode.yaml" export DSV4_MODEL_ID="${DSV4_MODEL_ID:-hf://deepseek-ai/DeepSeek-V4-Flash-0731}" -DATASET_PATH="${DATASET_ID:-/model/ljl/dataset/self-cognition.jsonl}" NPU_PER_NODE=16 NNODES=2 @@ -67,7 +66,6 @@ if [[ "${#VISIBLE_NPUS[@]}" -ne "$NPU_PER_NODE" ]]; then fi test -f "$CONFIG_PATH" -test -f "$DATASET_PATH" if [[ "$DSV4_MODEL_ID" == hf://* || "$DSV4_MODEL_ID" == ms://* ]]; then echo "Model will be downloaded through the configured Hub backend: $DSV4_MODEL_ID" diff --git a/cookbook/client/server/transformer/run_dsv4_4layer_ep_diagnostic.sh b/cookbook/client/server/transformer/run_dsv4_4layer_ep_diagnostic.sh deleted file mode 100755 index e1821f91f..000000000 --- a/cookbook/client/server/transformer/run_dsv4_4layer_ep_diagnostic.sh +++ /dev/null @@ -1,37 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail - -ROLE="${1:-}" -MODE="${2:-}" - -if [[ "$ROLE" != "head" && "$ROLE" != "worker" ]]; then - echo "Usage: $0 {head|worker} {no_ep|ep_loop|ep_gmm}" >&2 - exit 2 -fi -if [[ "$MODE" != "no_ep" && "$MODE" != "ep_loop" && "$MODE" != "ep_gmm" ]]; then - echo "Usage: $0 {head|worker} {no_ep|ep_loop|ep_gmm}" >&2 - exit 2 -fi - -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" - -case "$MODE" in - no_ep) - export TWINKLE_SERVER_CONFIG_PATH="$SCRIPT_DIR/server_config_dsv4_4layer_diag_no_ep.yaml" - export TWINKLE_EP_FORCE_LOOP=0 - ;; - ep_loop) - export TWINKLE_SERVER_CONFIG_PATH="$SCRIPT_DIR/server_config_dsv4_4layer_diag_ep.yaml" - export TWINKLE_EP_FORCE_LOOP=1 - ;; - ep_gmm) - export TWINKLE_SERVER_CONFIG_PATH="$SCRIPT_DIR/server_config_dsv4_4layer_diag_ep.yaml" - export TWINKLE_EP_FORCE_LOOP=0 - ;; -esac - -export TWINKLE_EP_DIAGNOSTICS="${TWINKLE_EP_DIAGNOSTICS:-1}" -export SKIP_DATASET_CHECK=1 - -echo "Starting DeepSeek-V4 four-layer diagnostic: role=$ROLE mode=$MODE" -exec "$SCRIPT_DIR/run_dsv4_0731_npu_2node_2npu.sh" "$ROLE" diff --git a/cookbook/client/server/transformer/server_config_dsv4_0731_npu.yaml b/cookbook/client/server/transformer/server_config_dsv4_0731_npu.yaml deleted file mode 100644 index 6d00f66eb..000000000 --- a/cookbook/client/server/transformer/server_config_dsv4_0731_npu.yaml +++ /dev/null @@ -1,116 +0,0 @@ -# DeepSeek-V4-Flash-0731 client/server Multi-LoRA SFT on one 2-NPU node. -# Both NPUs are assigned to the Transformers training model. A sampler is -# intentionally omitted because it would compete with training for devices. - -proxy_location: EveryNode - -http_options: - host: 0.0.0.0 - port: 8000 - -telemetry: - enabled: false - otlp_endpoint: http://localhost:4317 - -persistence: - mode: file - file_path: /tmp/twinkle_dsv4_0731_npu_state.json - -applications: - - name: server - route_prefix: /api/v1 - import_path: server - args: - server_config: - per_token_model_limit: 2 - supported_models: - - deepseek-v4-0731-local - deployments: - - name: TinkerCompatServer - max_ongoing_requests: 50 - autoscaling_config: - min_replicas: 1 - max_replicas: 1 - target_ongoing_requests: 32 - ray_actor_options: - num_cpus: 0.1 - runtime_env: - env_vars: - TWINKLE_FAIL_FAST: "1" - - - name: models-deepseek-v4-0731-local - route_prefix: /api/v1/model/deepseek-v4-0731-local - import_path: model - args: - backend: transformers - model_id: /nas/disk1/random-deepseek-v4-4b - nproc_per_node: 2 - max_loras: 2 - max_r: 8 - max_length: 2048 - mixed_precision: bf16 - strategy: native_fsdp - memory_efficient_init: true - target_modules: all-linear - fsdp_config: - reshard_after_forward: true - expert_parallel: - enabled: true - ep_size: 2 - router_dtype: fp32 - keep_router_logits: false - device_group: - name: model - ranks: 2 - device_type: npu - device_mesh: - device_type: npu - fsdp_size: 2 - dp_size: 1 - ep_size: 2 - queue_config: - rps_limit: 20 - tps_limit: 100000 - max_input_tokens: 4096 - queue_timeout: 3600 - execution_timeout: 3600 - adapter_config: - adapter_timeout: 3600 - adapter_max_lifetime: 86400 - deployments: - - name: ModelManagement - autoscaling_config: - min_replicas: 1 - max_replicas: 1 - target_ongoing_requests: 8 - ray_actor_options: - num_cpus: 0.1 - runtime_env: - env_vars: - TWINKLE_TRUST_REMOTE_CODE: "1" - TWINKLE_FAIL_FAST: "1" - - - name: processor - route_prefix: /api/v1/processor - import_path: processor - args: - ncpu_proc_per_node: 2 - device_group: - name: processor - ranks: 2 - device_type: CPU - device_mesh: - device_type: CPU - dp_size: 2 - deployments: - - name: ProcessorManagement - autoscaling_config: - min_replicas: 1 - max_replicas: 1 - target_ongoing_requests: 32 - ray_actor_options: - num_cpus: 0.1 - runtime_env: - env_vars: - TWINKLE_TRUST_REMOTE_CODE: "1" - TWINKLE_FAIL_FAST: "1" diff --git a/cookbook/client/server/transformer/server_config_dsv4_0731_npu_2node_2npu.yaml b/cookbook/client/server/transformer/server_config_dsv4_0731_npu_2node_2npu.yaml deleted file mode 100644 index f03121f8e..000000000 --- a/cookbook/client/server/transformer/server_config_dsv4_0731_npu_2node_2npu.yaml +++ /dev/null @@ -1,121 +0,0 @@ -# DeepSeek-V4-compatible Multi-LoRA SFT on two Ascend A3 nodes. -# Each node contributes 2 NPUs, giving a 4-rank Native FSDP2 + EP mesh. -# -# DSV4_MODEL_ID may be either: -# - hf://deepseek-ai/DeepSeek-V4-Flash-0731 -# - ms://namespace/model-name -# - the same absolute local model directory on both nodes - -proxy_location: EveryNode - -http_options: - host: 0.0.0.0 - port: 8000 - -telemetry: - enabled: false - otlp_endpoint: http://localhost:4317 - -persistence: - mode: memory - key_prefix: dsv4-0731-a3-2node-4npu - -applications: - - name: server - route_prefix: /api/v1 - import_path: server - args: - server_config: - per_token_model_limit: 2 - supported_models: - - deepseek-v4-0731-local - deployments: - - name: TinkerCompatServer - max_ongoing_requests: 50 - autoscaling_config: - min_replicas: 1 - max_replicas: 1 - target_ongoing_requests: 32 - ray_actor_options: - num_cpus: 0.1 - runtime_env: - env_vars: - TWINKLE_FAIL_FAST: "1" - - - name: models-deepseek-v4-0731-local - route_prefix: /api/v1/model/deepseek-v4-0731-local - import_path: model - args: - backend: transformers - # Resolved by OmegaConf when the head node reads this YAML. - model_id: ${oc.env:DSV4_MODEL_ID,hf://deepseek-ai/DeepSeek-V4-Flash-0731} - nproc_per_node: 2 - max_loras: 2 - max_r: 8 - max_length: 2048 - mixed_precision: bf16 - strategy: native_fsdp - memory_efficient_init: true - target_modules: all-linear - fsdp_config: - reshard_after_forward: true - expert_parallel: - enabled: true - ep_size: 4 - router_dtype: fp32 - keep_router_logits: false - device_group: - name: model - ranks: 4 - device_type: npu - device_mesh: - device_type: npu - fsdp_size: 4 - dp_size: 1 - ep_size: 4 - queue_config: - rps_limit: 20 - tps_limit: 100000 - max_input_tokens: 4096 - queue_timeout: 7200 - execution_timeout: 7200 - adapter_config: - adapter_timeout: 60 - adapter_max_lifetime: 86400 - deployments: - - name: ModelManagement - autoscaling_config: - min_replicas: 1 - max_replicas: 1 - target_ongoing_requests: 8 - ray_actor_options: - num_cpus: 0.1 - runtime_env: - env_vars: - TWINKLE_TRUST_REMOTE_CODE: "1" - TWINKLE_FAIL_FAST: "1" - - - name: processor - route_prefix: /api/v1/processor - import_path: processor - args: - ncpu_proc_per_node: 2 - device_group: - name: processor - ranks: 2 - device_type: CPU - device_mesh: - device_type: CPU - dp_size: 2 - deployments: - - name: ProcessorManagement - autoscaling_config: - min_replicas: 1 - max_replicas: 1 - target_ongoing_requests: 32 - ray_actor_options: - num_cpus: 0.1 - runtime_env: - env_vars: - TWINKLE_TRUST_REMOTE_CODE: "1" - TWINKLE_FAIL_FAST: "1" diff --git a/cookbook/client/server/transformer/server_config_dsv4_0731_npu_multinode.yaml b/cookbook/client/server/transformer/server_config_dsv4_0731_npu_multinode.yaml index c82324fad..e8ce7112e 100644 --- a/cookbook/client/server/transformer/server_config_dsv4_0731_npu_multinode.yaml +++ b/cookbook/client/server/transformer/server_config_dsv4_0731_npu_multinode.yaml @@ -30,7 +30,7 @@ applications: server_config: per_token_model_limit: 1 supported_models: - - deepseek-v4-0731-local + - DeepSeek-V4-Flash-0731 deployments: - name: TinkerCompatServer max_ongoing_requests: 50 @@ -44,8 +44,8 @@ applications: env_vars: TWINKLE_FAIL_FAST: "1" - - name: models-deepseek-v4-0731-local - route_prefix: /api/v1/model/deepseek-v4-0731-local + - name: models-DeepSeek-V4-Flash-0731 + route_prefix: /api/v1/model/DeepSeek-V4-Flash-0731 import_path: model args: backend: transformers @@ -54,8 +54,8 @@ applications: # Resolved from the launch environment on the head node. model_id: ${oc.env:DSV4_MODEL_ID,hf://deepseek-ai/DeepSeek-V4-Flash-0731} nproc_per_node: 16 - max_loras: 1 - max_r: 8 + max_loras: 3 + max_r: 32 max_length: 8192 mixed_precision: bf16 strategy: native_fsdp diff --git a/cookbook/client/server/transformer/server_config_dsv4_4layer_diag_ep.yaml b/cookbook/client/server/transformer/server_config_dsv4_4layer_diag_ep.yaml deleted file mode 100644 index cdbc82d76..000000000 --- a/cookbook/client/server/transformer/server_config_dsv4_4layer_diag_ep.yaml +++ /dev/null @@ -1,119 +0,0 @@ -# Four-layer DeepSeek-V4 diagnostic: FSDP2 + EP on two nodes, two NPUs per node. -# DSV4_MODEL_ID must point to a structurally complete four-layer BF16 checkpoint. - -proxy_location: EveryNode - -http_options: - host: 0.0.0.0 - port: 8000 - -telemetry: - enabled: false - otlp_endpoint: http://localhost:4317 - -persistence: - mode: memory - key_prefix: dsv4-4layer-diag-ep - -applications: - - name: server - route_prefix: /api/v1 - import_path: server - args: - server_config: - per_token_model_limit: 1 - supported_models: - - deepseek-v4-0731-local - deployments: - - name: TinkerCompatServer - max_ongoing_requests: 4 - autoscaling_config: - min_replicas: 1 - max_replicas: 1 - target_ongoing_requests: 1 - ray_actor_options: - num_cpus: 0.1 - runtime_env: - env_vars: - TWINKLE_FAIL_FAST: "1" - - - name: models-deepseek-v4-0731-local - route_prefix: /api/v1/model/deepseek-v4-0731-local - import_path: model - args: - backend: transformers - model_id: ${oc.env:DSV4_MODEL_ID} - nproc_per_node: 2 - max_loras: 1 - max_r: 8 - max_length: 512 - mixed_precision: bf16 - strategy: native_fsdp - memory_efficient_init: true - target_modules: all-linear - fsdp_config: - reshard_after_forward: true - expert_parallel: - enabled: true - ep_size: 4 - router_dtype: fp32 - keep_router_logits: false - device_group: - name: model - ranks: 4 - device_type: npu - device_mesh: - device_type: npu - fsdp_size: 4 - dp_size: 1 - ep_size: 4 - queue_config: - # Functional diagnostics issue several serial setup requests in quick - # succession; do not let the production-style limiter mask EP results. - rps_limit: 100 - tps_limit: 100000 - max_input_tokens: 512 - queue_timeout: 7200 - execution_timeout: 7200 - adapter_config: - adapter_timeout: 60 - adapter_max_lifetime: 7200 - deployments: - - name: ModelManagement - autoscaling_config: - min_replicas: 1 - max_replicas: 1 - target_ongoing_requests: 1 - ray_actor_options: - num_cpus: 0.1 - runtime_env: - env_vars: - TWINKLE_TRUST_REMOTE_CODE: "1" - TWINKLE_FAIL_FAST: "1" - TWINKLE_EP_FORCE_LOOP: ${oc.env:TWINKLE_EP_FORCE_LOOP,0} - TWINKLE_EP_DIAGNOSTICS: ${oc.env:TWINKLE_EP_DIAGNOSTICS,1} - - - name: processor - route_prefix: /api/v1/processor - import_path: processor - args: - ncpu_proc_per_node: 1 - device_group: - name: processor - ranks: 2 - device_type: CPU - device_mesh: - device_type: CPU - dp_size: 2 - deployments: - - name: ProcessorManagement - autoscaling_config: - min_replicas: 1 - max_replicas: 1 - target_ongoing_requests: 1 - ray_actor_options: - num_cpus: 0.1 - runtime_env: - env_vars: - TWINKLE_TRUST_REMOTE_CODE: "1" - TWINKLE_FAIL_FAST: "1" diff --git a/cookbook/client/server/transformer/server_config_dsv4_4layer_diag_no_ep.yaml b/cookbook/client/server/transformer/server_config_dsv4_4layer_diag_no_ep.yaml deleted file mode 100644 index 8be2c5951..000000000 --- a/cookbook/client/server/transformer/server_config_dsv4_4layer_diag_no_ep.yaml +++ /dev/null @@ -1,119 +0,0 @@ -# Four-layer DeepSeek-V4 diagnostic reference: FSDP2 without expert parallel. -# Apart from the EP settings, keep this aligned with server_config_dsv4_4layer_diag_ep.yaml. - -proxy_location: EveryNode - -http_options: - host: 0.0.0.0 - port: 8000 - -telemetry: - enabled: false - otlp_endpoint: http://localhost:4317 - -persistence: - mode: memory - key_prefix: dsv4-4layer-diag-no-ep - -applications: - - name: server - route_prefix: /api/v1 - import_path: server - args: - server_config: - per_token_model_limit: 1 - supported_models: - - deepseek-v4-0731-local - deployments: - - name: TinkerCompatServer - max_ongoing_requests: 4 - autoscaling_config: - min_replicas: 1 - max_replicas: 1 - target_ongoing_requests: 1 - ray_actor_options: - num_cpus: 0.1 - runtime_env: - env_vars: - TWINKLE_FAIL_FAST: "1" - - - name: models-deepseek-v4-0731-local - route_prefix: /api/v1/model/deepseek-v4-0731-local - import_path: model - args: - backend: transformers - model_id: ${oc.env:DSV4_MODEL_ID} - nproc_per_node: 2 - max_loras: 1 - max_r: 8 - max_length: 512 - mixed_precision: bf16 - strategy: native_fsdp - memory_efficient_init: true - target_modules: all-linear - fsdp_config: - reshard_after_forward: true - expert_parallel: - enabled: false - ep_size: 1 - router_dtype: fp32 - keep_router_logits: false - device_group: - name: model - ranks: 4 - device_type: npu - device_mesh: - device_type: npu - fsdp_size: 4 - dp_size: 1 - ep_size: 1 - queue_config: - # Functional diagnostics issue several serial setup requests in quick - # succession; do not let the production-style limiter mask EP results. - rps_limit: 100 - tps_limit: 100000 - max_input_tokens: 512 - queue_timeout: 7200 - execution_timeout: 7200 - adapter_config: - adapter_timeout: 60 - adapter_max_lifetime: 7200 - deployments: - - name: ModelManagement - autoscaling_config: - min_replicas: 1 - max_replicas: 1 - target_ongoing_requests: 1 - ray_actor_options: - num_cpus: 0.1 - runtime_env: - env_vars: - TWINKLE_TRUST_REMOTE_CODE: "1" - TWINKLE_FAIL_FAST: "1" - TWINKLE_EP_FORCE_LOOP: "0" - TWINKLE_EP_DIAGNOSTICS: ${oc.env:TWINKLE_EP_DIAGNOSTICS,1} - - - name: processor - route_prefix: /api/v1/processor - import_path: processor - args: - ncpu_proc_per_node: 1 - device_group: - name: processor - ranks: 2 - device_type: CPU - device_mesh: - device_type: CPU - dp_size: 2 - deployments: - - name: ProcessorManagement - autoscaling_config: - min_replicas: 1 - max_replicas: 1 - target_ongoing_requests: 1 - ray_actor_options: - num_cpus: 0.1 - runtime_env: - env_vars: - TWINKLE_TRUST_REMOTE_CODE: "1" - TWINKLE_FAIL_FAST: "1" diff --git a/cookbook/client/server/transformer/server_dsv4_0731_npu.py b/cookbook/client/server/transformer/server_dsv4_0731_npu.py deleted file mode 100644 index 7b19f24c3..000000000 --- a/cookbook/client/server/transformer/server_dsv4_0731_npu.py +++ /dev/null @@ -1,12 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Launch the two-NPU DeepSeek-V4-0731 Multi-LoRA server.""" -import os - -os.environ.setdefault('TWINKLE_TRUST_REMOTE_CODE', '1') - -from twinkle.server import launch_server # noqa: E402 - -file_dir = os.path.abspath(os.path.dirname(__file__)) -config_path = os.path.join(file_dir, 'server_config_dsv4_0731_npu.yaml') - -launch_server(config_path=config_path) diff --git a/cookbook/client/twinkle/demo_client_dataset.py b/cookbook/client/twinkle/demo_client_dataset.py new file mode 100644 index 000000000..56d4cdce9 --- /dev/null +++ b/cookbook/client/twinkle/demo_client_dataset.py @@ -0,0 +1,131 @@ +"""Twinkle C/S training demo with dataset processing on the client. + +The client loads, tokenizes, and batches the dataset. Only model operations are +sent to the Twinkle server. The input JSON/JSONL is expected to contain a +``messages`` column. +""" + +from peft import LoraConfig + +import twinkle +from twinkle import get_logger, init_twinkle_client +from twinkle.dataloader import DataLoader +from twinkle.dataset import Dataset, DatasetMeta +from twinkle_client.model import MultiLoraTransformersModel + +logger = get_logger() + +base_url = '{{ base_url }}' +api_key = '{{ api_key }}' +served_model_name = '{{ base_model }}' +model_path = '{{ model_local_path }}' +dataset_id = '{{ dataset_path }}' +epochs = {{epochs}} +max_length = {{max_length}} +truncation_strategy = '{{ truncation_strategy }}' +batch_size = 32 +lr = {{lr}} +grad_accumulation_steps = {{grad_accumulation_steps}} +train_id = '{{ train_id }}' + +template = 'DeepseekV4Template' +lora_target_modules = 'all-linear' +lora_rank = 16 +lora_alpha = 32 +lora_dropout = 0.0 + +train_config = { + 'base_model': served_model_name, + 'epochs': epochs, + 'max_length': max_length, + 'batch_size': batch_size, + 'lr': lr, + 'grad_accumulation_steps': grad_accumulation_steps, + 'lora_target_modules': lora_target_modules, + 'lora_rank': lora_rank, + 'lora_alpha': lora_alpha, + 'lora_dropout': lora_dropout, +} +logger.info(f'train config: {train_config}') + + +def _build_lora_config() -> LoraConfig: + return LoraConfig( + r=lora_rank, + lora_alpha=lora_alpha, + lora_dropout=lora_dropout, + target_modules='all-linear', + exclude_modules=['o_a_proj'], + target_parameters=[ + 'mlp.experts.gate_up_proj', + 'mlp.experts.down_proj', + ], + ) + + +def build_local_dataloader() -> DataLoader: + import os + if not os.path.exists(dataset_id): + raise FileNotFoundError(f'Client cannot access dataset: {dataset_id}') + if not os.path.exists(model_path): + raise FileNotFoundError(f'Client cannot access tokenizer/model directory: {model_path}') + + dataset = Dataset(dataset_meta=DatasetMeta(dataset_id=dataset_id)) + dataset.set_template( + template, + model_id=model_path, + max_length=max_length, + truncation_strategy=truncation_strategy, + ) + + # The input is already in {"messages": [...]} format. For raw + # query/response SelfCognition data, call dataset.map(...) before encode(). + dataset.encode(num_proc=8, load_from_cache_file=True) + return DataLoader( + dataset=dataset, + batch_size=batch_size, + drop_last=True, + shuffle=True, + num_workers=2, + ) + + +def train() -> None: + # This ensures twinkle.dataset and twinkle.dataloader execute in this process. + twinkle.initialize(mode='local') + + client = init_twinkle_client(base_url=base_url, api_key=api_key) + supported_models = [item.model_name for item in client.get_server_capabilities().supported_models] + if served_model_name not in supported_models: + raise RuntimeError(f'{served_model_name!r} is not served; available models: {supported_models}') + + dataloader = build_local_dataloader() + model = MultiLoraTransformersModel(model_id=served_model_name) + model.add_adapter_to_model( + 'default', + _build_lora_config(), + gradient_accumulation_steps=grad_accumulation_steps, + save_dir=None, + ) + # The model server automatically supplies its tokenizer_id here. + model.set_template(template) + model.set_processor('InputProcessor', padding_side='right') + model.set_loss('CrossEntropyLoss') + model.set_optimizer('Adam', lr=lr) + + for epoch in range(epochs): + logger.info('Starting epoch %s', epoch) + for step, batch in enumerate(dataloader, start=1): + model.forward_backward(inputs=batch) + model.clip_grad_and_step() + if step % grad_accumulation_steps == 0: + metric = model.calculate_metric(is_training=True) + logger.info('step=%s/%s metric=%s', step, len(dataloader), metric.result) + + checkpoint = model.save(name='twinkle-final', save_optimizer=True) + logger.info('Saved checkpoint: %s', checkpoint.twinkle_path) + client.close() + + +if __name__ == '__main__': + train() diff --git a/cookbook/client/twinkle/demo_server_dataset.py b/cookbook/client/twinkle/demo_server_dataset.py new file mode 100644 index 000000000..ef72e5957 --- /dev/null +++ b/cookbook/client/twinkle/demo_server_dataset.py @@ -0,0 +1,122 @@ +"""Twinkle C/S training demo with dataset processing on the server. + +Dataset, template encoding, and batching are executed by the remote Processor +service. Dataset and tokenizer/model paths must therefore be visible from every +Ray node on which a Processor Dataset actor may run. +""" + +from peft import LoraConfig + +from twinkle import get_logger, init_twinkle_client +from twinkle.dataset import DatasetMeta +from twinkle_client.dataloader import DataLoader +from twinkle_client.dataset import Dataset +from twinkle_client.model import MultiLoraTransformersModel + +logger = get_logger() + +base_url = '{{ base_url }}' +api_key = '{{ api_key }}' +served_model_name = '{{ base_model }}' +model_path = '{{ model_local_path }}' +dataset_id = '{{ dataset_path }}' +epochs = {{epochs}} +max_length = {{max_length}} +truncation_strategy = '{{ truncation_strategy }}' +batch_size = 32 +lr = {{lr}} +grad_accumulation_steps = {{grad_accumulation_steps}} +train_id = '{{ train_id }}' + +template = 'DeepseekV4Template' +lora_target_modules = 'all-linear' +lora_rank = 16 +lora_alpha = 32 +lora_dropout = 0.0 + +train_config = { + 'base_model': served_model_name, + 'epochs': epochs, + 'max_length': max_length, + 'batch_size': batch_size, + 'lr': lr, + 'grad_accumulation_steps': grad_accumulation_steps, + 'lora_target_modules': lora_target_modules, + 'lora_rank': lora_rank, + 'lora_alpha': lora_alpha, + 'lora_dropout': lora_dropout, +} +logger.info(f'train config: {train_config}') + + +def _build_lora_config() -> LoraConfig: + return LoraConfig( + r=lora_rank, + lora_alpha=lora_alpha, + lora_dropout=lora_dropout, + target_modules='all-linear', + exclude_modules=['o_a_proj'], + target_parameters=[ + 'mlp.experts.gate_up_proj', + 'mlp.experts.down_proj', + ], + ) + + +def build_remote_dataloader() -> DataLoader: + # Do not use os.path.exists() here: it would only check the client machine. + dataset = Dataset(dataset_meta=DatasetMeta(dataset_id=dataset_id)) + dataset.set_template( + template, + model_id=model_path, + max_length=max_length, + truncation_strategy=truncation_strategy, + ) + + # The input is already in {"messages": [...]} format. For raw + # query/response SelfCognition data, call dataset.map(...) before encode(). + dataset.encode(timeout=3600, num_proc=8, load_from_cache_file=True) + return DataLoader( + dataset=dataset, + batch_size=batch_size, + min_batch_size=batch_size, + drop_last=True, + shuffle=True, + ) + + +def train() -> None: + client = init_twinkle_client(base_url=base_url, api_key=api_key) + supported_models = [item.model_name for item in client.get_server_capabilities().supported_models] + if served_model_name not in supported_models: + raise RuntimeError(f'{served_model_name!r} is not served; available models: {supported_models}') + + dataloader = build_remote_dataloader() + model = MultiLoraTransformersModel(model_id=served_model_name) + model.add_adapter_to_model( + 'default', + _build_lora_config(), + gradient_accumulation_steps=grad_accumulation_steps, + save_dir=None, + ) + model.set_template(template) + model.set_processor('InputProcessor', padding_side='right') + model.set_loss('CrossEntropyLoss') + model.set_optimizer('Adam', lr=lr) + + for epoch in range(epochs): + logger.info('Starting epoch %s', epoch) + for step, batch in enumerate(dataloader, start=1): + model.forward_backward(inputs=batch) + model.clip_grad_and_step() + if step % grad_accumulation_steps == 0: + metric = model.calculate_metric(is_training=True) + logger.info('step=%s/%s metric=%s', step, len(dataloader), metric.result) + + checkpoint = model.save(name='twinkle-final', save_optimizer=True) + logger.info('Saved checkpoint: %s', checkpoint.twinkle_path) + client.close() + + +if __name__ == '__main__': + train() diff --git a/cookbook/client/twinkle/dsv4_multi_lora_self_cognition_sft.py b/cookbook/client/twinkle/dsv4_multi_lora_self_cognition_sft.py index e52b58aa4..46be4b27a 100644 --- a/cookbook/client/twinkle/dsv4_multi_lora_self_cognition_sft.py +++ b/cookbook/client/twinkle/dsv4_multi_lora_self_cognition_sft.py @@ -1,198 +1,148 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""Client-side self-cognition SFT for a DeepSeek-V4-Flash-0731 server. +"""DeepSeek-V4 self-cognition SFT with client-side dataset processing.""" -Start the DeepSeek-V4 multi-node server first. This client trains one LoRA -tenant and updates only the routed-expert gate/up/down parameters. -""" -import math import os +import twinkle from peft import LoraConfig - from twinkle import get_logger, init_twinkle_client -from twinkle.dataset import DatasetMeta -from twinkle_client.dataloader import DataLoader -from twinkle_client.dataset import Dataset +from twinkle.dataloader import DataLoader +from twinkle.dataset import Dataset, DatasetMeta from twinkle_client.model import MultiLoraTransformersModel + logger = get_logger() -SERVER_URL = os.environ.get('TWINKLE_SERVER_URL', 'http://localhost:8000') -SERVER_TOKEN = os.environ.get('TWINKLE_SERVER_TOKEN', 'EMPTY_TOKEN') -SERVED_MODEL_NAME = os.environ.get('TWINKLE_MODEL_ID', 'deepseek-v4-0731-local') -MODEL_PATH = os.environ.get('DSV4_MODEL_ID', '/nas/disk1/random-deepseek-v4-4b') -DATASET_ID = os.environ.get('DATASET_ID') -OUTPUT_DIR = os.environ.get('OUTPUT_DIR', '/tmp/twinkle_dsv4_0731_self_cognition') -MODEL_NAME = os.environ.get('SELF_COGNITION_MODEL_NAME', 'twinkle模型') -MODEL_AUTHOR = os.environ.get('SELF_COGNITION_MODEL_AUTHOR', 'ModelScope社区') - -ADAPTER_NAME = os.environ.get('ADAPTER_NAME', 'tenant_a').strip() -BATCH_SIZE = int(os.environ.get('BATCH_SIZE', '32')) -GRAD_ACCUM_STEPS = int(os.environ.get('GRAD_ACCUM_STEPS', '1')) -NUM_EPOCHS = int(os.environ.get('NUM_EPOCHS', '3')) -MAX_STEPS = int(os.environ.get('MAX_STEPS', '0')) -MAX_LENGTH = int(os.environ.get('MAX_LENGTH', '8192')) -TRUNCATION_STRATEGY = os.environ.get('TRUNCATION_STRATEGY', 'delete') -LR = float(os.environ.get('LR', '1e-4')) -LORA_R = int(os.environ.get('LORA_R', '8')) -LORA_ALPHA = int(os.environ.get('LORA_ALPHA', '32')) -ROUTED_EXPERT_TARGET_PARAMETERS = [ - 'mlp.experts.gate_up_proj', - 'mlp.experts.down_proj', -] - - -def _assert_finite_output(value, path='result') -> None: - if isinstance(value, dict): - for key, item in value.items(): - _assert_finite_output(item, f'{path}.{key}') - elif isinstance(value, (list, tuple)): - for index, item in enumerate(value): - _assert_finite_output(item, f'{path}[{index}]') - elif isinstance(value, float) and not math.isfinite(value): - raise RuntimeError(f'Inference smoke test produced a non-finite value at {path}: {value}') - - -def _build_dataset() -> Dataset: - if not DATASET_ID: - raise ValueError('Set DATASET_ID to a local self-cognition JSON or JSONL file.') - if DATASET_ID.startswith(('hf://', 'ms://')): - raise ValueError(f'DATASET_ID must be local for this recipe, got: {DATASET_ID}') - if not os.path.exists(DATASET_ID): - raise FileNotFoundError(f'Local self-cognition dataset not found: {DATASET_ID}') - - dataset = Dataset(dataset_meta=DatasetMeta(DATASET_ID)) - dataset.set_template( - 'DeepseekV4Template', - model_id=MODEL_PATH, - max_length=MAX_LENGTH, - truncation_strategy=TRUNCATION_STRATEGY, - ) - dataset.map( - 'SelfCognitionProcessor', - init_args={ - 'model_name': MODEL_NAME, - 'model_author': MODEL_AUTHOR, - }, - ) - dataset.encode() - return dataset +base_url = os.environ.get('TWINKLE_SERVER_URL', 'http://localhost:8000') +api_key = os.environ.get('TWINKLE_SERVER_TOKEN', 'EMPTY_TOKEN') +served_model_name = os.environ.get('TWINKLE_MODEL_ID', 'DeepSeek-V4-Flash-0731') +model_path = os.environ.get('MODEL_LOCAL_PATH') or os.environ.get('DSV4_MODEL_ID', '') +dataset_id = os.environ.get('DATASET_PATH') or os.environ.get('DATASET_ID', '') +epochs = int(os.environ.get('EPOCHS') or os.environ.get('NUM_EPOCHS', '3')) +max_length = int(os.environ.get('MAX_LENGTH', '8192')) +truncation_strategy = os.environ.get('TRUNCATION_STRATEGY', 'delete') +batch_size = int(os.environ.get('BATCH_SIZE', '32')) +lr = float(os.environ.get('LR', '1e-4')) +grad_accumulation_steps = int( + os.environ.get('GRAD_ACCUMULATION_STEPS') or os.environ.get('GRAD_ACCUM_STEPS', '1') +) +train_id = os.environ.get('TRAIN_ID') or os.environ.get('ADAPTER_NAME', 'tenant_a') + +template = 'DeepseekV4Template' +lora_target_modules = 'all-linear' +lora_rank = 16 +lora_alpha = 32 +lora_dropout = 0.0 + +train_config = { + 'base_model': served_model_name, + 'epochs': epochs, + 'max_length': max_length, + 'batch_size': batch_size, + 'lr': lr, + 'grad_accumulation_steps': grad_accumulation_steps, + 'lora_target_modules': lora_target_modules, + 'lora_rank': lora_rank, + 'lora_alpha': lora_alpha, + 'lora_dropout': lora_dropout, +} +logger.info(f'train config: {train_config}') def _build_lora_config() -> LoraConfig: return LoraConfig( - r=LORA_R, - lora_alpha=LORA_ALPHA, - lora_dropout=0.0, - target_modules=None, - target_parameters=ROUTED_EXPERT_TARGET_PARAMETERS, - bias='none', + r=lora_rank, + lora_alpha=lora_alpha, + lora_dropout=lora_dropout, + target_modules='all-linear', + exclude_modules=['o_a_proj'], + target_parameters=[ + 'mlp.experts.gate_up_proj', + 'mlp.experts.down_proj', + ], ) -def train() -> None: - client = init_twinkle_client(base_url=SERVER_URL, api_key=SERVER_TOKEN) - supported_models = [item.model_name for item in client.get_server_capabilities().supported_models] - if SERVED_MODEL_NAME not in supported_models: - raise RuntimeError(f'{SERVED_MODEL_NAME!r} is not served; available models: {supported_models}') - if not ADAPTER_NAME: - raise ValueError('ADAPTER_NAME must not be empty.') - if BATCH_SIZE < 32 or BATCH_SIZE % 32 != 0: - raise ValueError('BATCH_SIZE must be at least 32 and divisible by 32 for the 32-rank FSDP model.') - if NUM_EPOCHS <= 0: - raise ValueError('NUM_EPOCHS must be greater than zero.') - if MAX_STEPS < 0: - raise ValueError('MAX_STEPS must be zero (unlimited) or a positive integer.') - - dataset = _build_dataset() - dataloader = DataLoader( - dataset=dataset, - batch_size=BATCH_SIZE, - drop_last=True, - shuffle=True, - ) - model = MultiLoraTransformersModel(model_id=SERVED_MODEL_NAME) - save_dir = os.path.join(OUTPUT_DIR, ADAPTER_NAME) - os.makedirs(save_dir, exist_ok=True) - model.add_adapter_to_model( - ADAPTER_NAME, - _build_lora_config(), - gradient_accumulation_steps=GRAD_ACCUM_STEPS, - save_dir=save_dir, - ) - model.set_template( - 'DeepseekV4Template', - max_length=MAX_LENGTH, - truncation_strategy=TRUNCATION_STRATEGY, - ) - model.set_processor('InputProcessor', padding_side='right') - model.set_loss('CrossEntropyLoss') - model.set_optimizer('AdamW', lr=LR, foreach=False) - - completed_steps = 0 - stop_training = False - for epoch in range(NUM_EPOCHS): - logger.info('Starting epoch %s/%s', epoch + 1, NUM_EPOCHS) - for batch in dataloader: - if MAX_STEPS > 0 and completed_steps >= MAX_STEPS: - stop_training = True - break - model.forward_backward(inputs=batch, gradient_accumulation_steps=GRAD_ACCUM_STEPS) - model.clip_grad_and_step(max_grad_norm=1.0, gradient_accumulation_steps=GRAD_ACCUM_STEPS) - completed_steps += 1 - - if completed_steps % GRAD_ACCUM_STEPS == 0: - metric = model.calculate_metric(is_training=True) - logger.info( - 'adapter=%s epoch=%s/%s micro_step=%s metric=%s', - ADAPTER_NAME, - epoch + 1, - NUM_EPOCHS, - completed_steps, - metric.result, - ) - if stop_training: - break - - if completed_steps == 0: - raise RuntimeError( - f'No full batch was produced: dataset is smaller than BATCH_SIZE={BATCH_SIZE}. ' - 'Use a smaller valid global batch or provide more data.' +def build_local_dataloader() -> DataLoader: + if not dataset_id: + raise ValueError('Set DATASET_PATH or DATASET_ID to the client-local dataset path.') + if not model_path: + raise ValueError( + 'Set MODEL_LOCAL_PATH or DSV4_MODEL_ID to the client-local tokenizer/model directory.' ) + if not os.path.exists(dataset_id): + raise FileNotFoundError(f'Client cannot access dataset: {dataset_id}') + if not os.path.exists(model_path): + raise FileNotFoundError(f'Client cannot access tokenizer/model directory: {model_path}') - checkpoint = model.save( - name=f'dsv4-0731-{ADAPTER_NAME}-self-cognition-final', - save_optimizer=True, - consumed_train_samples=dataloader.get_state()['consumed_train_samples'], + dataset = Dataset(dataset_meta=DatasetMeta(dataset_id=dataset_id)) + dataset.set_template( + template, + model_id=model_path, + max_length=max_length, + truncation_strategy=truncation_strategy, ) - logger.info('Saved adapter %s: %s', ADAPTER_NAME, checkpoint.twinkle_path) - - eval_loader = DataLoader( + dataset.map( + 'SelfCognitionProcessor', + init_args={ + 'model_name': 'twinkle模型', + 'model_author': 'twinkle团队', + }, + ) + dataset.encode(num_proc=8, load_from_cache_file=True) + return DataLoader( dataset=dataset, - batch_size=BATCH_SIZE, + batch_size=batch_size, drop_last=True, - shuffle=False, - ) - eval_batch = next(iter(eval_loader)) - lora_result = model.forward_only(inputs=eval_batch, disable_lora=False).result - lora_loss = model.calculate_loss().result - base_result = model.forward_only(inputs=eval_batch, disable_lora=True).result - base_loss = model.calculate_loss().result - _assert_finite_output(lora_result, f'{ADAPTER_NAME}.lora') - _assert_finite_output(base_result, f'{ADAPTER_NAME}.base') - _assert_finite_output(lora_loss, f'{ADAPTER_NAME}.lora_loss') - _assert_finite_output(base_loss, f'{ADAPTER_NAME}.base_loss') - logger.warning( - 'Inference result: adapter=%s lora_loss=%.6f base_loss=%.6f loss_delta=%.6f ' - 'lora_keys=%s base_keys=%s', - ADAPTER_NAME, - lora_loss, - base_loss, - lora_loss - base_loss, - list(lora_result) if isinstance(lora_result, dict) else type(lora_result).__name__, - list(base_result) if isinstance(base_result, dict) else type(base_result).__name__, + shuffle=True, + num_workers=2, ) +def train() -> None: + # Dataset loading, mapping, tokenization, and batching all run locally. + twinkle.initialize(mode='local') + + client = init_twinkle_client(base_url=base_url, api_key=api_key) + try: + supported_models = [item.model_name for item in client.get_server_capabilities().supported_models] + if served_model_name not in supported_models: + raise RuntimeError( + f'{served_model_name!r} is not served; available models: {supported_models}' + ) + + dataloader = build_local_dataloader() + model = MultiLoraTransformersModel(model_id=served_model_name) + model.add_adapter_to_model( + train_id, + _build_lora_config(), + gradient_accumulation_steps=grad_accumulation_steps, + save_dir=None, + ) + model.set_template(template) + model.set_processor('InputProcessor', padding_side='right') + model.set_loss('CrossEntropyLoss') + model.set_optimizer('Adam', lr=lr) + + for epoch in range(epochs): + logger.info('Starting epoch %s/%s', epoch + 1, epochs) + for step, batch in enumerate(dataloader, start=1): + model.forward_backward( + inputs=batch, + gradient_accumulation_steps=grad_accumulation_steps, + ) + model.clip_grad_and_step( + gradient_accumulation_steps=grad_accumulation_steps, + ) + if step % grad_accumulation_steps == 0: + metric = model.calculate_metric(is_training=True) + logger.info('step=%s/%s metric=%s', step, len(dataloader), metric.result) + + checkpoint = model.save(name=f'{train_id}-final', save_optimizer=True) + logger.info('Saved checkpoint: %s', checkpoint.twinkle_path) + finally: + client.close() + + if __name__ == '__main__': train() diff --git a/cookbook/client/twinkle/dsv4_multi_lora_sft.py b/cookbook/client/twinkle/dsv4_multi_lora_sft.py deleted file mode 100644 index 24438a526..000000000 --- a/cookbook/client/twinkle/dsv4_multi_lora_sft.py +++ /dev/null @@ -1,203 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Client-side GSM8K LoRA SFT for a DeepSeek-V4-Flash-0731 server. - -Start the DeepSeek-V4 multi-node server first. This client trains one LoRA -tenant on the EP/FSDP model. -""" -import math -import os - -from peft import LoraConfig - -from twinkle import get_logger, init_twinkle_client -from twinkle.dataset import DatasetMeta -from twinkle_client.dataloader import DataLoader -from twinkle_client.dataset import Dataset -from twinkle_client.model import MultiLoraTransformersModel - -logger = get_logger() - -SERVER_URL = os.environ.get('TWINKLE_SERVER_URL', 'http://localhost:8000') -SERVER_TOKEN = os.environ.get('TWINKLE_SERVER_TOKEN', 'EMPTY_TOKEN') -SERVED_MODEL_NAME = os.environ.get('TWINKLE_MODEL_ID', 'deepseek-v4-0731-local') -MODEL_PATH = os.environ.get('DSV4_MODEL_ID', '/nas/disk1/random-deepseek-v4-4b') -DATASET_ID = os.environ.get('DATASET_ID') -DATASET_SUBSET = os.environ.get('DATASET_SUBSET', 'default') -DATASET_SPLIT = os.environ.get('DATASET_SPLIT', 'train') -OUTPUT_DIR = os.environ.get('OUTPUT_DIR', '/tmp/twinkle_dsv4_0731_multi_lora') - -ADAPTER_NAMES = tuple( - name.strip() for name in os.environ.get('ADAPTER_NAMES', 'tenant_a').split(',') if name.strip()) -BATCH_SIZE = int(os.environ.get('BATCH_SIZE', '32')) -GRAD_ACCUM_STEPS = int(os.environ.get('GRAD_ACCUM_STEPS', '4')) -MAX_STEPS = int(os.environ.get('MAX_STEPS', '10')) -MAX_LENGTH = int(os.environ.get('MAX_LENGTH', '8192')) -TRUNCATION_STRATEGY = os.environ.get('TRUNCATION_STRATEGY', 'delete') -LR = float(os.environ.get('LR', '1e-4')) -LORA_R = int(os.environ.get('LORA_R', '8')) -LORA_ALPHA = int(os.environ.get('LORA_ALPHA', '32')) -ROUTED_EXPERT_TARGET_PARAMETERS = [ - 'mlp.experts.gate_up_proj', - 'mlp.experts.down_proj', -] - - -def _assert_finite_output(value, path='result') -> None: - """Fail the inference smoke test when a returned numeric value is NaN/Inf.""" - if isinstance(value, dict): - for key, item in value.items(): - _assert_finite_output(item, f'{path}.{key}') - elif isinstance(value, (list, tuple)): - for index, item in enumerate(value): - _assert_finite_output(item, f'{path}[{index}]') - elif isinstance(value, float) and not math.isfinite(value): - raise RuntimeError(f'Inference smoke test produced a non-finite value at {path}: {value}') - - -def _build_dataset() -> Dataset: - if not DATASET_ID: - raise ValueError('Set DATASET_ID to a local GSM8K JSON/JSONL file or directory.') - if DATASET_ID.startswith(('hf://', 'ms://')): - raise ValueError(f'DATASET_ID must be local for this recipe, got: {DATASET_ID}') - if not os.path.exists(DATASET_ID): - raise FileNotFoundError(f'Local GSM8K dataset not found: {DATASET_ID}') - dataset = Dataset(dataset_meta=DatasetMeta( - DATASET_ID, - subset_name=DATASET_SUBSET, - split=DATASET_SPLIT, - )) - dataset.set_template( - 'DeepseekV4Template', - model_id=MODEL_PATH, - max_length=MAX_LENGTH, - truncation_strategy=TRUNCATION_STRATEGY, - ) - dataset.map( - 'GSM8KProcessor', - init_args={ - 'system': 'Solve the math problem step by step and put the final answer in \\boxed{}.', - # SFT needs the reference solution as the assistant target. - 'add_assistant': True, - }, - ) - dataset.encode() - return dataset - - -def _build_lora_config() -> LoraConfig: - return LoraConfig( - r=LORA_R, - lora_alpha=LORA_ALPHA, - lora_dropout=0.0, - # Train only the fused 3D routed-expert parameters. Do not install - # module LoRA on attention, router, shared experts, or the LM head. - target_modules=None, - target_parameters=ROUTED_EXPERT_TARGET_PARAMETERS, - bias='none', - ) - - -def train() -> None: - client = init_twinkle_client(base_url=SERVER_URL, api_key=SERVER_TOKEN) - supported_models = [item.model_name for item in client.get_server_capabilities().supported_models] - if SERVED_MODEL_NAME not in supported_models: - raise RuntimeError(f'{SERVED_MODEL_NAME!r} is not served; available models: {supported_models}') - if len(ADAPTER_NAMES) != 1: - raise ValueError('This recipe requires exactly one ADAPTER_NAMES entry.') - if BATCH_SIZE < 32 or BATCH_SIZE % 32 != 0: - raise ValueError('BATCH_SIZE must be at least 32 and divisible by 32 for the 32-rank FSDP model.') - - dataset = _build_dataset() - # Drop the final undersized batch because 32 FSDP ranks require a full - # global batch here. GSM8K is large enough for the configured run, so each - # adapter traverses at most one shuffled epoch. - dataloaders = { - name: DataLoader(dataset=dataset, batch_size=BATCH_SIZE, drop_last=True, shuffle=True) - for name in ADAPTER_NAMES - } - models = {name: MultiLoraTransformersModel(model_id=SERVED_MODEL_NAME) for name in ADAPTER_NAMES} - lora_config = _build_lora_config() - - # Register every tenant before set_optimizer(). The first optimizer call - # materializes and EP/FSDP-shards all preallocated Multi-LoRA slots. - for name, model in models.items(): - save_dir = os.path.join(OUTPUT_DIR, name) - os.makedirs(save_dir, exist_ok=True) - model.add_adapter_to_model( - name, - lora_config, - gradient_accumulation_steps=GRAD_ACCUM_STEPS, - save_dir=save_dir, - ) - model.set_template( - 'DeepseekV4Template', - max_length=MAX_LENGTH, - truncation_strategy=TRUNCATION_STRATEGY, - ) - model.set_processor('InputProcessor', padding_side='right') - model.set_loss('CrossEntropyLoss') - - for model in models.values(): - model.set_optimizer('AdamW', lr=LR, foreach=False) - - iterators = {name: iter(loader) for name, loader in dataloaders.items()} - completed_steps = {name: 0 for name in ADAPTER_NAMES} - active_names = set(ADAPTER_NAMES) - - while active_names and any(step < MAX_STEPS for step in completed_steps.values()): - for name in ADAPTER_NAMES: - if name not in active_names or completed_steps[name] >= MAX_STEPS: - continue - try: - batch = next(iterators[name]) - except StopIteration: - active_names.remove(name) - continue - - model = models[name] - model.forward_backward(inputs=batch, gradient_accumulation_steps=GRAD_ACCUM_STEPS) - model.clip_grad_and_step(max_grad_norm=1.0, gradient_accumulation_steps=GRAD_ACCUM_STEPS) - completed_steps[name] += 1 - - if completed_steps[name] % GRAD_ACCUM_STEPS == 0: - metric = model.calculate_metric(is_training=True) - logger.info('adapter=%s micro_step=%s metric=%s', name, completed_steps[name], metric.result) - - for name, model in models.items(): - checkpoint = model.save( - name=f'dsv4-0731-{name}-final', - save_optimizer=True, - consumed_train_samples=dataloaders[name].get_state()['consumed_train_samples'], - ) - logger.info('Saved adapter %s: %s', name, checkpoint.twinkle_path) - - # Run a real no-grad forward pass after saving. Compare the trained LoRA - # path with the base-model path so both can be verified on the same input. - # Saving first preserves the expensive training result if this check fails. - eval_loader = DataLoader(dataset=dataset, batch_size=BATCH_SIZE, drop_last=True, shuffle=False) - eval_batch = next(iter(eval_loader)) - for name, model in models.items(): - lora_result = model.forward_only(inputs=eval_batch, disable_lora=False).result - lora_loss = model.calculate_loss().result - base_result = model.forward_only(inputs=eval_batch, disable_lora=True).result - base_loss = model.calculate_loss().result - _assert_finite_output(lora_result, f'{name}.lora') - _assert_finite_output(base_result, f'{name}.base') - _assert_finite_output(lora_loss, f'{name}.lora_loss') - _assert_finite_output(base_loss, f'{name}.base_loss') - # Use WARNING so the inference result remains visible even when the - # deployment intentionally suppresses INFO logs to save local storage. - logger.warning( - 'Inference result: adapter=%s lora_loss=%.6f base_loss=%.6f loss_delta=%.6f ' - 'lora_keys=%s base_keys=%s', - name, - lora_loss, - base_loss, - lora_loss - base_loss, - list(lora_result) if isinstance(lora_result, dict) else type(lora_result).__name__, - list(base_result) if isinstance(base_result, dict) else type(base_result).__name__, - ) - - -if __name__ == '__main__': - train() diff --git a/cookbook/client/twinkle/run_dsv4_0731_npu_2node_2npu_client.sh b/cookbook/client/twinkle/run_dsv4_0731_npu_2node_2npu_client.sh deleted file mode 100755 index 5cce5dfce..000000000 --- a/cookbook/client/twinkle/run_dsv4_0731_npu_2node_2npu_client.sh +++ /dev/null @@ -1,33 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail - -# Run exactly one client process. The default server address points at the -# two-node cluster head configured in run_dsv4_0731_npu_2node_2npu.sh. - -HEAD_IP="${HEAD_IP:-172.61.10.111}" -OUTPUT_DIR="${OUTPUT_DIR:-/shared/twinkle_output/dsv4-0731-a3-2node-4npu}" - -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" -PROJECT_DIR="$(cd -- "$SCRIPT_DIR/../../.." && pwd)" -cd "$PROJECT_DIR" - -export PYTHONPATH="$PROJECT_DIR/src${PYTHONPATH:+:$PYTHONPATH}" -export TWINKLE_SERVER_URL="${TWINKLE_SERVER_URL:-http://$HEAD_IP:8000}" -export TWINKLE_SERVER_TOKEN="${TWINKLE_SERVER_TOKEN:-EMPTY_TOKEN}" -export TWINKLE_MODEL_ID="${TWINKLE_MODEL_ID:-deepseek-v4-0731-local}" -export DSV4_MODEL_ID="${DSV4_MODEL_ID:-hf://deepseek-ai/DeepSeek-V4-Flash-0731}" -export DATASET_ID="${DATASET_ID:-/model/ljl/dataset/self-cognition.jsonl}" -export ADAPTER_NAMES="${ADAPTER_NAMES:-tenant_a,tenant_b}" -export OUTPUT_DIR - -# Four global FSDP ranks: one sample per rank by default. -export BATCH_SIZE="${BATCH_SIZE:-4}" -export GRAD_ACCUM_STEPS="${GRAD_ACCUM_STEPS:-4}" -export MAX_STEPS="${MAX_STEPS:-10}" -export MAX_LENGTH="${MAX_LENGTH:-2048}" -export LORA_R="${LORA_R:-8}" -export LORA_ALPHA="${LORA_ALPHA:-32}" -export LR="${LR:-1e-4}" - -mkdir -p "$OUTPUT_DIR" -exec python3 "$SCRIPT_DIR/dsv4_multi_lora_sft.py" diff --git a/cookbook/client/twinkle/run_dsv4_0731_npu_client.sh b/cookbook/client/twinkle/run_dsv4_0731_npu_client.sh deleted file mode 100755 index 60068b705..000000000 --- a/cookbook/client/twinkle/run_dsv4_0731_npu_client.sh +++ /dev/null @@ -1,26 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail - -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" -PROJECT_DIR="$(cd -- "$SCRIPT_DIR/../../.." && pwd)" -cd "$PROJECT_DIR" - -export PYTHONPATH="$PROJECT_DIR/src${PYTHONPATH:+:$PYTHONPATH}" -export TWINKLE_SERVER_URL="${TWINKLE_SERVER_URL:-http://127.0.0.1:8000}" -export TWINKLE_SERVER_TOKEN="${TWINKLE_SERVER_TOKEN:-EMPTY_TOKEN}" -export TWINKLE_MODEL_ID="${TWINKLE_MODEL_ID:-deepseek-v4-0731-local}" -export DSV4_MODEL_ID="${DSV4_MODEL_ID:-/nas/disk1/random-deepseek-v4-4b}" -export DATASET_ID="${DATASET_ID:-/model/ljl/dataset/self-cognition.jsonl}" -export ADAPTER_NAMES="${ADAPTER_NAMES:-tenant_a,tenant_b}" -export BATCH_SIZE="${BATCH_SIZE:-2}" -export GRAD_ACCUM_STEPS="${GRAD_ACCUM_STEPS:-4}" -export MAX_STEPS="${MAX_STEPS:-10}" -export MAX_LENGTH="${MAX_LENGTH:-2048}" -export LORA_R="${LORA_R:-8}" -export LORA_ALPHA="${LORA_ALPHA:-32}" -export LR="${LR:-1e-4}" -export OUTPUT_DIR="${OUTPUT_DIR:-/tmp/twinkle_dsv4_0731_npu_multi_lora}" - -test -f "$DATASET_ID" - -exec python3 "$SCRIPT_DIR/dsv4_multi_lora_sft.py" diff --git a/cookbook/client/twinkle/run_dsv4_0731_npu_multinode_client.sh b/cookbook/client/twinkle/run_dsv4_0731_npu_multinode_client.sh deleted file mode 100755 index e56542d48..000000000 --- a/cookbook/client/twinkle/run_dsv4_0731_npu_multinode_client.sh +++ /dev/null @@ -1,47 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail - -# The client is an HTTP orchestrator. Run exactly one client process; it does -# not join Ray and does not need NPU devices. -# -# Example: -# HEAD_IP=10.0.0.10 \ -# OUTPUT_DIR=/shared/twinkle_output/dsv4-0731-a3-32npu \ -# bash run_dsv4_0731_npu_multinode_client.sh - -: "${HEAD_IP:?Set HEAD_IP to the Twinkle/Ray head node IP}" -: "${OUTPUT_DIR:?Set OUTPUT_DIR to a shared absolute path mounted on both server nodes}" - -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" -PROJECT_DIR="$(cd -- "$SCRIPT_DIR/../../.." && pwd)" -cd "$PROJECT_DIR" - -export PYTHONPATH="$PROJECT_DIR/src${PYTHONPATH:+:$PYTHONPATH}" -export TWINKLE_SERVER_URL="${TWINKLE_SERVER_URL:-http://$HEAD_IP:8000}" -export TWINKLE_SERVER_TOKEN="${TWINKLE_SERVER_TOKEN:-EMPTY_TOKEN}" -export TWINKLE_MODEL_ID="${TWINKLE_MODEL_ID:-deepseek-v4-0731-local}" -export ADAPTER_NAMES="${ADAPTER_NAMES:-tenant_a}" -# The template/encoding side must resolve the same checkpoint as the server. -export DSV4_MODEL_ID="${DSV4_MODEL_ID:-hf://deepseek-ai/DeepSeek-V4-Flash-0731}" -: "${DATASET_ID:?Set DATASET_ID to a local GSM8K JSON/JSONL file or directory}" -if [[ "$DATASET_ID" == hf://* || "$DATASET_ID" == ms://* ]]; then - echo "DATASET_ID must be local for this recipe, got: $DATASET_ID" >&2 - exit 1 -fi -test -e "$DATASET_ID" -export DATASET_ID -export DATASET_SUBSET="${DATASET_SUBSET:-default}" -export DATASET_SPLIT="${DATASET_SPLIT:-train}" -export OUTPUT_DIR - -# forward_backward dispatches one slice to each of the 32 FSDP ranks. -export BATCH_SIZE="${BATCH_SIZE:-32}" -export GRAD_ACCUM_STEPS="${GRAD_ACCUM_STEPS:-4}" -export MAX_STEPS="${MAX_STEPS:-10}" -export MAX_LENGTH="${MAX_LENGTH:-8192}" -export TRUNCATION_STRATEGY="${TRUNCATION_STRATEGY:-delete}" -export LORA_R="${LORA_R:-8}" -export LORA_ALPHA="${LORA_ALPHA:-32}" - -mkdir -p "$OUTPUT_DIR" -exec python3 cookbook/client/twinkle/dsv4_multi_lora_sft.py diff --git a/scripts/process_self_cognition_dataset.py b/scripts/process_self_cognition_dataset.py new file mode 100644 index 000000000..3ad8edec3 --- /dev/null +++ b/scripts/process_self_cognition_dataset.py @@ -0,0 +1,86 @@ +#!/usr/bin/env python3 +"""Convert a SelfCognition JSON/JSONL dataset to chat messages JSONL.""" + +import argparse +import json +from pathlib import Path +from typing import Any, Iterator + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description='Convert query/response SelfCognition samples into messages format.') + parser.add_argument('input', type=Path, help='Input .json or .jsonl file') + parser.add_argument('output', type=Path, help='Output .jsonl file') + parser.add_argument('--model-name', default='twinkle模型', help='Value used for {{NAME}}') + parser.add_argument('--model-author', default='twinkle团队', help='Value used for {{AUTHOR}}') + return parser.parse_args() + + +def read_samples(path: Path) -> Iterator[dict[str, Any]]: + if path.suffix.lower() == '.jsonl': + with path.open('r', encoding='utf-8') as source: + for line_number, line in enumerate(source, start=1): + if not line.strip(): + continue + value = json.loads(line) + if not isinstance(value, dict): + raise ValueError(f'Line {line_number} must be a JSON object') + yield value + return + + with path.open('r', encoding='utf-8') as source: + value = json.load(source) + + if isinstance(value, dict): + value = [value] + if not isinstance(value, list): + raise ValueError('A JSON input must contain an object or a list of objects') + for index, sample in enumerate(value): + if not isinstance(sample, dict): + raise ValueError(f'Item {index} must be a JSON object') + yield sample + + +def replace_placeholders(text: str, model_name: str, model_author: str) -> str: + return text.replace('{{NAME}}', model_name).replace('{{AUTHOR}}', model_author) + + +def convert_sample(sample: dict[str, Any], model_name: str, model_author: str) -> dict[str, Any]: + if 'query' not in sample or 'response' not in sample: + raise ValueError("Each sample must contain both 'query' and 'response'") + if not isinstance(sample['query'], str) or not isinstance(sample['response'], str): + raise ValueError("The 'query' and 'response' fields must be strings") + + return { + 'messages': [ + { + 'role': 'system', + 'content': 'You are a helpful assistant.' + }, + { + 'role': 'user', + 'content': replace_placeholders(sample['query'], model_name, model_author), + }, + { + 'role': 'assistant', + 'content': replace_placeholders(sample['response'], model_name, model_author), + }, + ] + } + + +def main() -> None: + args = parse_args() + args.output.parent.mkdir(parents=True, exist_ok=True) + + count = 0 + with args.output.open('w', encoding='utf-8') as destination: + for count, sample in enumerate(read_samples(args.input), start=1): + converted = convert_sample(sample, args.model_name, args.model_author) + destination.write(json.dumps(converted, ensure_ascii=False) + '\n') + + print(f'Converted {count} samples to {args.output}') + + +if __name__ == '__main__': + main() diff --git a/src/twinkle/kernel/ops/ep/__init__.py b/src/twinkle/kernel/ops/ep/__init__.py index 556d1cff7..c011d3ced 100644 --- a/src/twinkle/kernel/ops/ep/__init__.py +++ b/src/twinkle/kernel/ops/ep/__init__.py @@ -11,7 +11,6 @@ """ from __future__ import annotations -import os import torch from abc import ABC, abstractmethod from torch import nn @@ -52,11 +51,6 @@ def forward( _IMPLS: list[EpExpertsGmm] | None = None _PATH_LOGGED = False _WARN_LOGGED = False -_FORCE_LOOP_LOGGED = False - - -def _env_flag(name: str) -> bool: - return os.environ.get(name, '').strip().lower() in {'1', 'true', 'yes', 'on'} def _get_impls() -> list[EpExpertsGmm]: @@ -65,14 +59,6 @@ def _get_impls() -> list[EpExpertsGmm]: Each backend module is imported defensively: platforms lacking its dependencies (e.g. no torch_npu) simply skip that backend. """ - global _FORCE_LOOP_LOGGED - if _env_flag('TWINKLE_EP_FORCE_LOOP'): - from .loop import LoopEpExpertsGmm - if not _FORCE_LOOP_LOGGED: - logger.warning('EP experts compute: TWINKLE_EP_FORCE_LOOP=1; forcing the per-expert F.linear loop.') - _FORCE_LOOP_LOGGED = True - return [LoopEpExpertsGmm()] - global _IMPLS if _IMPLS is None: _IMPLS = [] diff --git a/src/twinkle/kernel/ops/moe/npu.py b/src/twinkle/kernel/ops/moe/npu.py index 8fcbfa694..fb32f2813 100644 --- a/src/twinkle/kernel/ops/moe/npu.py +++ b/src/twinkle/kernel/ops/moe/npu.py @@ -63,9 +63,8 @@ def _normalize_packed_expert_weights(module, input_dtype, hidden_dim): gate_up_proj = module.gate_up_proj.to(input_dtype) down_proj = module.down_proj.to(input_dtype) if gate_up_proj.ndim != 3 or down_proj.ndim != 3: - raise RuntimeError( - 'Packed expert weights must be 3D: ' - f'gate_up_proj={tuple(gate_up_proj.shape)}, down_proj={tuple(down_proj.shape)}.') + raise RuntimeError('Packed expert weights must be 3D: ' + f'gate_up_proj={tuple(gate_up_proj.shape)}, down_proj={tuple(down_proj.shape)}.') # torch.nn.functional.linear stores weights as [out_features, in_features], # while torch_npu.npu_grouped_matmul consumes [in_features, out_features]. @@ -76,11 +75,10 @@ def _normalize_packed_expert_weights(module, input_dtype, hidden_dim): linear_layout = gate_up_proj.shape[2] == hidden_dim and down_proj.shape[1] == hidden_dim grouped_mm_layout = gate_up_proj.shape[1] == hidden_dim and down_proj.shape[2] == hidden_dim if linear_layout == grouped_mm_layout: - raise RuntimeError( - 'Unable to determine packed expert weight layout: ' - f'gate_up_proj={tuple(gate_up_proj.shape)}, down_proj={tuple(down_proj.shape)}, ' - f'hidden_dim={hidden_dim}. Expected either Transformers/F.linear ' - '[E, out, in] tensors or grouped-matmul [E, in, out] tensors.') + raise RuntimeError('Unable to determine packed expert weight layout: ' + f'gate_up_proj={tuple(gate_up_proj.shape)}, down_proj={tuple(down_proj.shape)}, ' + f'hidden_dim={hidden_dim}. Expected either Transformers/F.linear ' + '[E, out, in] tensors or grouped-matmul [E, in, out] tensors.') if linear_layout: gate_up_weight = gate_up_proj.transpose(1, 2) diff --git a/src/twinkle/model/multi_lora.py b/src/twinkle/model/multi_lora.py index 1104bfa84..7048d1c73 100644 --- a/src/twinkle/model/multi_lora.py +++ b/src/twinkle/model/multi_lora.py @@ -173,15 +173,10 @@ def patch_target_parameters(self, module, target_parameters): self.target_parameter_manager.patch(module, target_parameters) @contextmanager - def adapter(self, - tenant_adapter_name: str, - disable_lora: bool = False, - cache_target_parameters: bool = True): + def adapter(self, tenant_adapter_name: str, disable_lora: bool = False, cache_target_parameters: bool = True): self.activate_adapter(tenant_adapter_name) with self.target_parameter_manager.adapter( - tenant_adapter_name, - disable_lora=disable_lora, - cache=cache_target_parameters): + tenant_adapter_name, disable_lora=disable_lora, cache=cache_target_parameters): if disable_lora: # Temporarily disable all adapters while keeping optimizer_group active with self._disable_lora_context(tenant_adapter_name): diff --git a/src/twinkle/model/transformers/moe/expert_parallel.py b/src/twinkle/model/transformers/moe/expert_parallel.py index 1ec1f3361..efc6e4b9e 100644 --- a/src/twinkle/model/transformers/moe/expert_parallel.py +++ b/src/twinkle/model/transformers/moe/expert_parallel.py @@ -2,30 +2,16 @@ from __future__ import annotations import inspect -import os import torch import torch.distributed as dist from dataclasses import dataclass from torch import nn from typing import Any, Dict, Iterable, List, Optional, Tuple -from twinkle import get_logger from twinkle.kernel.ops import ep_forward from twinkle.model.transformers.moe.ep_utils import preprocess, token_pre_all2all, tokens_post_all2all from twinkle.utils import DeviceMesh -logger = get_logger() - - -def _ep_diagnostics_enabled() -> bool: - return os.environ.get('TWINKLE_EP_DIAGNOSTICS', '').strip().lower() in {'1', 'true', 'yes', 'on'} - - -def _split_sizes_to_list(split_sizes: list[int] | torch.Tensor) -> list[int]: - if isinstance(split_sizes, torch.Tensor): - return [int(value) for value in split_sizes.detach().cpu().tolist()] - return [int(value) for value in split_sizes] - @dataclass class ExpertParallelConfig: @@ -299,34 +285,6 @@ def forward(hidden_states: torch.Tensor, *args, **kwargs): num_global_sum_tokens_per_local_expert, ) = preprocess(expert_mask, num_experts, ep_group) - log_diagnostics = _ep_diagnostics_enabled() and not getattr(block, '_ep_diagnostics_logged', False) - if log_diagnostics: - input_splits_list = _split_sizes_to_list(input_splits) - output_splits_list = _split_sizes_to_list(output_splits) - expected_assignments = hidden_states_2d.shape[0] * top_k - actual_assignments = sum(input_splits_list) - if actual_assignments != expected_assignments: - raise RuntimeError( - f'EP routing assignment mismatch for {block_name}: ' - f'input_splits={actual_assignments}, expected={expected_assignments}.') - selected_preview = selected_experts[:min(8, selected_experts.shape[0])].detach().cpu().tolist() - weights_preview = routing_weights[:min(8, routing_weights.shape[0])].float().detach().cpu().tolist() - logger.warning( - '[EP_DIAG] rank=%s ep_rank=%s block=%s expert_range=[%s,%s) tokens=%s top_k=%s ' - 'input_splits=%s output_splits=%s selected_experts=%s routing_weights=%s', - dist.get_rank(), - block._ep_rank, - block_name, - block._ep_local_start, - block._ep_local_end, - hidden_states_2d.shape[0], - top_k, - input_splits_list, - output_splits_list, - selected_preview, - weights_preview, - ) - # 2. token_pre_all2all: permute → all_to_all → sort_chunks ( global_permuted_hidden_states, @@ -382,21 +340,6 @@ def forward(hidden_states: torch.Tensor, *args, **kwargs): if len(orig_shape) == 3: final_hidden = final_hidden.view(batch_size, seq_len, hidden_dim) - if log_diagnostics: - flat_output = final_hidden.detach().reshape(-1) - preview = flat_output[:min(8, flat_output.numel())].float().cpu().tolist() - finite = bool(torch.isfinite(final_hidden).all().item()) - logger.warning( - '[EP_DIAG] rank=%s ep_rank=%s block=%s output_shape=%s output_finite=%s output_preview=%s', - dist.get_rank(), - block._ep_rank, - block_name, - tuple(final_hidden.shape), - finite, - preview, - ) - block._ep_diagnostics_logged = True - if cfg.keep_router_logits and returns_router_logits: return final_hidden, router_logits return final_hidden diff --git a/src/twinkle/model/transformers/multi_lora_transformers.py b/src/twinkle/model/transformers/multi_lora_transformers.py index 54b828ef9..0dce75c70 100644 --- a/src/twinkle/model/transformers/multi_lora_transformers.py +++ b/src/twinkle/model/transformers/multi_lora_transformers.py @@ -207,10 +207,7 @@ def generate(self, # Generation invokes many forwards inside one context. Do not retain a # parametrized expert weight across FSDP reshard boundaries; recompute # the routed-expert LoRA delta when each decoder forward accesses it. - with self.multi_adapter.adapter( - adapter_name, - disable_lora=disable_lora, - cache_target_parameters=False): + with self.multi_adapter.adapter(adapter_name, disable_lora=disable_lora, cache_target_parameters=False): return super().generate( inputs=inputs, adapter_name=adapter_name, diff --git a/src/twinkle/model/transformers/strategy/native_fsdp.py b/src/twinkle/model/transformers/strategy/native_fsdp.py index 049ec2a59..c44adf2a1 100644 --- a/src/twinkle/model/transformers/strategy/native_fsdp.py +++ b/src/twinkle/model/transformers/strategy/native_fsdp.py @@ -8,15 +8,13 @@ from torch.distributed.fsdp import fully_shard from typing import TYPE_CHECKING, Any, Dict, List, Literal, Mapping, Optional, Set -from twinkle.utils import DeviceMesh, Platform, get_logger, torch_util +from twinkle.utils import DeviceMesh, Platform, torch_util from twinkle.utils.torch_utils import snapshot_state_dict_to_cpu from .load_context import fsdp_pretrained_load_context if TYPE_CHECKING: from torch.distributed.fsdp import MixedPrecisionPolicy -logger = get_logger() - LORA_STATE_KEY_MARKERS = ('lora_A', 'lora_B', 'lora_embedding') PEFT_BASE_PREFIX = 'base_model.model.' PEFT_BASE_LAYER_SEGMENT = 'base_layer' @@ -84,8 +82,7 @@ def capture_pre_ep_state_if_needed(self, model, *, enable_ep: bool) -> None: # second full-model-sized allocation on every node-local source rank. # The retained tensors are immutable until wrap_model() finishes # broadcasting the rank-local shards. - self.set_rank0_pre_ep_full_state_dict( - snapshot_state_dict_to_cpu(model.state_dict()) if is_source_rank else {}) + self.set_rank0_pre_ep_full_state_dict(snapshot_state_dict_to_cpu(model.state_dict()) if is_source_rank else {}) self._pre_ep_state_captured = True def can_reuse_pre_ep_tensor_storage(self) -> bool: @@ -1028,9 +1025,6 @@ def _broadcast_sharded_state_dict( rank_to_ep_rank = rank_to_ep_rank or {} adapter_source_sd = adapter_source_sd or {} adapter_full_sd = adapter_full_sd or {} - diagnostics_enabled = os.environ.get('TWINKLE_EP_DIAGNOSTICS', '').strip().lower() in { - '1', 'true', 'yes', 'on' - } source_metadata = None source_keys = None adapter_metadata = None @@ -1144,43 +1138,14 @@ def _scatter_ep_expert_tensor(param_name, full_tensor, sharded_param): raise RuntimeError(f"EP expert parameter '{param_name}' expects {num_experts} experts, " f'but source state has shape {tuple(full_tensor.shape)}. ' 'Rank0 must capture the full pre-EP state_dict before apply_expert_parallel().') - if diagnostics_enabled: - logger.warning( - '[EP_DIAG] rank=%s local_source=%s param=%s full_shape=%s source_preview=%s', - rank, - local_source_rank, - param_name, - tuple(full_tensor.shape), - _diagnostic_tensor_preview(full_tensor), - ) local_tensor = _scatter_ep_tensor_from_source( full_tensor, local_tensor, shard_dim=0, shard_size=experts_per_rank, ) - if diagnostics_enabled: - ep_rank = rank_to_ep_rank[rank] - start = ep_rank * experts_per_rank - logger.warning( - '[EP_DIAG] rank=%s ep_rank=%s param=%s expert_range=[%s,%s) local_shape=%s local_preview=%s', - rank, - ep_rank, - param_name, - start, - start + experts_per_rank, - tuple(local_tensor.shape), - _diagnostic_tensor_preview(local_tensor), - ) return local_tensor - def _diagnostic_tensor_preview(tensor: torch.Tensor) -> List[float]: - flat = tensor.detach().reshape(-1) - if flat.numel() == 0: - return [] - indices = sorted({0, flat.numel() // 3, (2 * flat.numel()) // 3, flat.numel() - 1}) - return flat[indices].float().cpu().tolist() - def _scatter_ep_tensor_from_source(full_tensor, local_tensor, *, shard_dim: int, shard_size: int): if is_source_rank: if full_tensor is None: diff --git a/tests/kernel/ops/test_ep_dispatch.py b/tests/kernel/ops/test_ep_dispatch.py index 3fde6031f..e89f73a71 100644 --- a/tests/kernel/ops/test_ep_dispatch.py +++ b/tests/kernel/ops/test_ep_dispatch.py @@ -1,7 +1,6 @@ import torch from twinkle.kernel.ops import ep as ep_ops -from twinkle.kernel.ops.ep import loop as loop_ops class _Backend(ep_ops.EpExpertsGmm): @@ -19,15 +18,6 @@ def forward(self, experts_mod, permuted_tokens, num_global_sum_tokens_per_local_ def test_ep_forward_uses_accelerated_backend_by_default(monkeypatch): - monkeypatch.delenv('TWINKLE_EP_FORCE_LOOP', raising=False) monkeypatch.setattr(ep_ops, '_IMPLS', [_Backend('accelerated', 1.0), _Backend('loop', 2.0, fallback=True)]) result = ep_ops.ep_forward(None, torch.zeros(2, 3), torch.tensor([2]), 1) assert torch.equal(result, torch.ones(2, 3)) - - -def test_ep_forward_can_force_reference_loop(monkeypatch): - monkeypatch.setenv('TWINKLE_EP_FORCE_LOOP', '1') - monkeypatch.setattr(ep_ops, '_IMPLS', [_Backend('accelerated', 1.0), _Backend('loop', 2.0, fallback=True)]) - monkeypatch.setattr(loop_ops, 'LoopEpExpertsGmm', lambda: _Backend('loop', 2.0, fallback=True)) - result = ep_ops.ep_forward(None, torch.zeros(2, 3), torch.tensor([2]), 1) - assert torch.equal(result, torch.full((2, 3), 2.0)) diff --git a/tests/kernel/ops/test_moe.py b/tests/kernel/ops/test_moe.py index e77bc9b6a..26c5692b2 100644 --- a/tests/kernel/ops/test_moe.py +++ b/tests/kernel/ops/test_moe.py @@ -1,5 +1,3 @@ -import importlib - import pytest import torch from torch import nn @@ -63,13 +61,3 @@ def test_normalize_packed_expert_weights_rejects_inconsistent_layout(): module = _PackedExperts(torch.randn(2, 7, 9), torch.randn(2, 5, 6)) with pytest.raises(RuntimeError, match='Unable to determine packed expert weight layout'): _normalize_packed_expert_weights(module, torch.float32, hidden_dim=8) - - -def test_ep_force_loop_environment_switch(monkeypatch): - ep_ops = importlib.import_module('twinkle.kernel.ops.ep') - - monkeypatch.setenv('TWINKLE_EP_FORCE_LOOP', '1') - implementations = ep_ops._get_impls() - - assert len(implementations) == 1 - assert implementations[0].name == 'per-expert loop'