From f49314895c8ce9276eead3d18948a9fcd256d3af Mon Sep 17 00:00:00 2001 From: zxc-phy Date: Fri, 31 Jul 2026 11:29:55 +0800 Subject: [PATCH 1/2] fix the bug of training error --- eval/configs/assets.json | 2 +- eval/experiments/ablation/run.py | 8 +- eval/tests/test_protocols.py | 35 +++- metis/checkpoint_utils.py | 112 +++++++++- metis/weight_utils.py | 348 +++++++++++++++++++++++++------ scripts/run_train.sh | 2 + scripts/train.sh | 11 +- train/dataset.py | 11 +- train/run_train.py | 135 ++++++++---- train/train_utils.py | 31 ++- train/trainer.py | 38 ++++ 11 files changed, 602 insertions(+), 131 deletions(-) diff --git a/eval/configs/assets.json b/eval/configs/assets.json index f9dac2f..0f95444 100644 --- a/eval/configs/assets.json +++ b/eval/configs/assets.json @@ -1,6 +1,6 @@ { "schema_version": 1, - "note": "Public model IDs may be replaced by local mirrors. Relative paths resolve from the Metis_dev repository root.", + "note": "Public model IDs may be replaced by local mirrors. Relative paths resolve from the repository root.", "assets": { "qwen3_5_4b": "Qwen/Qwen3.5-4B", "qwen3_5_9b": "Qwen/Qwen3.5-9B", diff --git a/eval/experiments/ablation/run.py b/eval/experiments/ablation/run.py index c4d8920..559dbdd 100644 --- a/eval/experiments/ablation/run.py +++ b/eval/experiments/ablation/run.py @@ -13,6 +13,7 @@ import shutil import subprocess import sys +import tempfile from collections import defaultdict from pathlib import Path from typing import Any, Iterable @@ -21,8 +22,9 @@ REPO_ROOT = Path(__file__).resolve().parents[3] -METIS_DEV_ROOT = Path(os.environ.get("METIS_MODEL_REPO_ROOT", str(REPO_ROOT))) +MODEL_REPO_ROOT = Path(os.environ.get("METIS_MODEL_REPO_ROOT", str(REPO_ROOT))) DEFAULT_MATRIX = Path(__file__).resolve().parent / "configs" / "ablation_matrix.json" +DEFAULT_JUDGE_LOCK = Path(tempfile.gettempdir()) / "metis_eval_judge32.lock" READY_STATUSES = {"ready", "reference_ready"} REQUIRED_CHECKPOINT_FILES = { "config.json", @@ -562,7 +564,7 @@ def run_raw_stage(args: argparse.Namespace, matrix: dict[str, Any]) -> dict[str, "matrix": str(args.matrix), "matrix_sha256": sha256(args.matrix), "git": git_snapshot(), - "metis_dev_git": git_snapshot(METIS_DEV_ROOT), + "model_repo_git": git_snapshot(MODEL_REPO_ROOT), "python": args.python, "checkpoints": selected, "benchmarks": benchmarks, @@ -915,7 +917,7 @@ def build_parser() -> argparse.ArgumentParser: parser.add_argument("--judge-concurrency", type=int, default=0) parser.add_argument("--judge-base-url", default="https://api.openai.com") parser.add_argument("--api-key-env", default="OPENAI_API_KEY") - parser.add_argument("--judge-lock", type=Path, default=Path("/tmp/metis_eval_judge32.lock")) + parser.add_argument("--judge-lock", type=Path, default=DEFAULT_JUDGE_LOCK) parser.add_argument("--judge-max-attempts", type=int, default=18) parser.add_argument("--judge-retry-sleep", type=float, default=2.0) parser.add_argument("--judge-retry-backoff", type=float, default=1.5) diff --git a/eval/tests/test_protocols.py b/eval/tests/test_protocols.py index c84739c..5abcdbf 100644 --- a/eval/tests/test_protocols.py +++ b/eval/tests/test_protocols.py @@ -1,6 +1,7 @@ from __future__ import annotations import json +import re import subprocess import sys import tempfile @@ -160,15 +161,31 @@ def test_ablation_checkpoints_use_the_shared_asset_registry() -> None: assert {item["id"] for item in matrix["checkpoints"]} <= set(assets) -def test_release_text_has_no_private_server_paths_or_usernames() -> None: +def test_release_text_has_no_private_paths_or_development_repo_names() -> None: root = Path(__file__).resolve().parents[2] - forbidden = ( - "/" + "mnt/afs/", + forbidden_literals = ( + "/" + "mnt/", "/" + "home/", + "/" + "root/", "/" + "Users/", + "/" + "tmp/", "api-" + "int.", ) - suffixes = {".py", ".json", ".md", ".toml", ".yml", ".txt"} + forbidden_repo_names = ( + "Metis_" + "dev", + "Metis_" + "open", + ) + suffixes = { + ".cff", + ".json", + ".md", + ".py", + ".sh", + ".toml", + ".txt", + ".yaml", + ".yml", + } candidates = subprocess.run( [ "git", @@ -177,7 +194,6 @@ def test_release_text_has_no_private_server_paths_or_usernames() -> None: "--cached", "--others", "--exclude-standard", - "eval", ], cwd=root, check=True, @@ -189,7 +205,14 @@ def test_release_text_has_no_private_server_paths_or_usernames() -> None: path = root / relative.decode() if path.is_file() and path.suffix in suffixes: text = path.read_text(encoding="utf-8", errors="replace").lower() - assert not any(value.lower() in text for value in forbidden), path + assert not any(value.lower() in text for value in forbidden_literals), path + assert not any( + re.search( + rf"(? None: diff --git a/metis/checkpoint_utils.py b/metis/checkpoint_utils.py index 38650b6..dfa3fa7 100644 --- a/metis/checkpoint_utils.py +++ b/metis/checkpoint_utils.py @@ -14,11 +14,17 @@ from typing import Any import torch +from safetensors import safe_open from safetensors.torch import load_file, save_file from .configuration_metis import MetisConfig from .modeling_metis import MetisForCausalLM -from .weight_utils import load_metis_from_backbone +from .weight_utils import ( + _gather_parameters_for_update, + _is_zero3_parameter, + is_deepspeed_zero3_enabled, + load_metis_from_backbone, +) logger = logging.getLogger(__name__) @@ -58,12 +64,29 @@ def _safe_dtype_name(tensor: torch.Tensor) -> str: def _trainable_state_dict(model, state_dict: dict[str, torch.Tensor] | None = None) -> dict[str, torch.Tensor]: unwrapped = _unwrap_model(model) - source = state_dict if state_dict is not None else unwrapped.state_dict() - trainable_names = [name for name, param in unwrapped.named_parameters() if param.requires_grad] + trainable_parameters = [ + (name, param) + for name, param in unwrapped.named_parameters() + if param.requires_grad + ] + + # A normal state_dict from a ZeRO-3 model contains only local + # placeholders. Gather one trainable tensor at a time instead of asking + # Accelerate/DeepSpeed to consolidate the complete frozen 27B model. + if state_dict is None and any( + _is_zero3_parameter(param) for _, param in trainable_parameters + ): + delta: dict[str, torch.Tensor] = {} + for name, parameter in trainable_parameters: + with _gather_parameters_for_update([parameter]) as should_collect: + if should_collect: + delta[name] = parameter.detach().cpu().clone() + return delta + source = state_dict if state_dict is not None else unwrapped.state_dict() delta: dict[str, torch.Tensor] = {} missing: list[str] = [] - for name in trainable_names: + for name, _ in trainable_parameters: tensor = source.get(name) if tensor is None: tensor = source.get(f"module.{name}") @@ -89,17 +112,30 @@ def save_metis_delta_checkpoint( tokenizer=None, base_model_path: str = "", state_dict: dict[str, torch.Tensor] | None = None, + is_main_process: bool = True, ) -> None: """Save a compact checkpoint containing only trainable Metis weights.""" output = Path(output_dir) - output.mkdir(parents=True, exist_ok=True) unwrapped = _unwrap_model(model) + uses_zero3_collectives = any( + _is_zero3_parameter(parameter) + for parameter in unwrapped.parameters() + ) if base_model_path: meta = dict(getattr(unwrapped.config, "backbone_meta", {}) or {}) meta["backbone_path"] = base_model_path unwrapped.config.backbone_meta = meta + # All ZeRO-3 ranks must participate in the per-parameter gathers above. + # Only rank 0 receives tensors and writes files. + delta = _trainable_state_dict(unwrapped, state_dict=state_dict) + if not is_main_process: + if uses_zero3_collectives: + torch.distributed.barrier() + return + + output.mkdir(parents=True, exist_ok=True) unwrapped.config.save_pretrained(output) generation_config = getattr(unwrapped, "generation_config", None) if generation_config is not None: @@ -107,7 +143,6 @@ def save_metis_delta_checkpoint( if tokenizer is not None: tokenizer.save_pretrained(output) - delta = _trainable_state_dict(unwrapped, state_dict=state_dict) save_file(delta, str(output / DELTA_WEIGHTS_NAME)) manifest = { @@ -133,6 +168,8 @@ def save_metis_delta_checkpoint( len(delta), manifest["num_parameters"] / 1_000_000, ) + if uses_zero3_collectives: + torch.distributed.barrier() def load_delta_manifest(checkpoint_path: str | os.PathLike[str]) -> dict[str, Any]: @@ -238,6 +275,65 @@ def load_metis_delta_into_model( raise ValueError(f"Not a Metis delta checkpoint: {checkpoint}") unwrapped = _unwrap_model(model) + if is_deepspeed_zero3_enabled(): + manifest = load_delta_manifest(checkpoint) + manifest_tensors = manifest.get("tensors") or {} + parameter_map = dict(unwrapped.named_parameters()) + file_path = checkpoint / DELTA_WEIGHTS_NAME + + with safe_open(str(file_path), framework="pt", device="cpu") as tensors: + file_names = set(tensors.keys()) + manifest_names = set(manifest_tensors) + if file_names != manifest_names: + missing = sorted(manifest_names - file_names) + unexpected = sorted(file_names - manifest_names) + raise RuntimeError( + "Delta manifest and safetensors file do not match: " + f"missing={missing[:10]}, unexpected={unexpected[:10]}" + ) + + unknown = sorted(file_names - parameter_map.keys()) + if unknown: + raise RuntimeError( + f"Delta checkpoint has {len(unknown)} unexpected keys: " + + ", ".join(unknown[:10]) + ) + + # Validate every logical shape before entering a collective. This + # prevents one rank from raising mid-gather while its peers wait. + for name in sorted(file_names): + parameter = parameter_map[name] + logical_shape = tuple( + getattr(parameter, "ds_shape", parameter.shape) + ) + checkpoint_shape = tuple(manifest_tensors[name]["shape"]) + if logical_shape != checkpoint_shape: + raise RuntimeError( + f"Delta tensor shape mismatch for {name}: " + f"model={logical_shape}, checkpoint={checkpoint_shape}" + ) + + # Materialize one trainable tensor at a time. Rank 0 reads it, + # then GatheredParameters broadcasts and re-partitions the update. + for name in sorted(file_names): + parameter = parameter_map[name] + with _gather_parameters_for_update([parameter]) as should_update: + if should_update: + value = tensors.get_tensor(name) + parameter.copy_( + value.to( + device=parameter.device, + dtype=parameter.dtype, + ) + ) + + logger.info( + "Loaded Metis delta weights into ZeRO-3 shards from %s (%d tensors)", + checkpoint, + len(file_names), + ) + return + delta = load_file(str(checkpoint / DELTA_WEIGHTS_NAME), device="cpu") incompatible = unwrapped.load_state_dict(delta, strict=False) if incompatible.unexpected_keys: @@ -267,10 +363,14 @@ def load_metis_model_from_checkpoint( dtype=dtype, ) load_metis_delta_into_model(model, checkpoint) + if is_deepspeed_zero3_enabled(): + return model return model.to(device=device, dtype=dtype) if is_full_checkpoint(checkpoint): model = MetisForCausalLM.from_pretrained(checkpoint, dtype=dtype) + if is_deepspeed_zero3_enabled(): + return model return model.to(device=device, dtype=dtype) raise ValueError( diff --git a/metis/weight_utils.py b/metis/weight_utils.py index f6ac30b..4180a60 100644 --- a/metis/weight_utils.py +++ b/metis/weight_utils.py @@ -2,7 +2,10 @@ from __future__ import annotations +from contextlib import contextmanager + import torch +from torch import nn from transformers import AutoModelForCausalLM, AutoTokenizer from .configuration_metis import MetisConfig @@ -16,6 +19,170 @@ } +def is_deepspeed_zero3_enabled() -> bool: + """Return whether Transformers currently has an active ZeRO-3 config. + + ``TrainingArguments`` must be constructed before calling the model loader + so Transformers can register the DeepSpeed config used by + ``PreTrainedModel._from_config`` and ``from_pretrained``. + """ + try: + from transformers.integrations.deepspeed import ( + is_deepspeed_zero3_enabled as transformers_zero3_enabled, + ) + except ImportError: + return False + return bool(transformers_zero3_enabled()) + + +def _is_zero3_parameter(parameter: nn.Parameter) -> bool: + return ( + hasattr(parameter, "ds_id") + and getattr(parameter, "ds_tensor", None) is not None + and hasattr(parameter, "ds_numel") + ) + + +@contextmanager +def _gather_parameters_for_update(parameters): + """Gather only the listed ZeRO-3 parameters and repartition on exit. + + All ranks enter the context, but only rank 0 mutates the materialized + tensors. ``modifier_rank=0`` then broadcasts the updated values before + DeepSpeed re-partitions them. + """ + unique: list[nn.Parameter] = [] + seen: set[int] = set() + for parameter in parameters: + if parameter is None or id(parameter) in seen: + continue + seen.add(id(parameter)) + unique.append(parameter) + + zero_parameters = [parameter for parameter in unique if _is_zero3_parameter(parameter)] + if not zero_parameters: + yield True + return + + import deepspeed + + with deepspeed.zero.GatheredParameters(zero_parameters, modifier_rank=0): + yield torch.distributed.get_rank() == 0 + + +def _copy_materialized_module_state( + target: nn.Module, + source: nn.Module, + *, + should_update: bool, +) -> bool: + """Copy a small, already-materialized module without creating state dicts.""" + target_parameters = dict(target.named_parameters()) + source_parameters = dict(source.named_parameters()) + target_buffers = dict(target.named_buffers()) + source_buffers = dict(source.named_buffers()) + if target_parameters.keys() != source_parameters.keys(): + return False + if target_buffers.keys() != source_buffers.keys(): + return False + + for name, target_parameter in target_parameters.items(): + source_parameter = source_parameters[name] + if target_parameter.shape != source_parameter.shape: + return False + if should_update: + target_parameter.copy_( + source_parameter.to( + device=target_parameter.device, + dtype=target_parameter.dtype, + ) + ) + + for name, target_buffer in target_buffers.items(): + source_buffer = source_buffers[name] + if target_buffer.shape != source_buffer.shape: + return False + if should_update: + target_buffer.copy_( + source_buffer.to(device=target_buffer.device, dtype=target_buffer.dtype) + ) + return True + + +@torch.no_grad() +def _copy_module_state(target: nn.Module, source: nn.Module) -> None: + """Copy module state, directly shard-to-shard when both models use ZeRO-3.""" + target_parameters = dict(target.named_parameters()) + source_parameters = dict(source.named_parameters()) + target_buffers = dict(target.named_buffers()) + source_buffers = dict(source.named_buffers()) + + missing = sorted(source_parameters.keys() - target_parameters.keys()) + unexpected = sorted(target_parameters.keys() - source_parameters.keys()) + missing_buffers = sorted(source_buffers.keys() - target_buffers.keys()) + unexpected_buffers = sorted(target_buffers.keys() - source_buffers.keys()) + if missing or unexpected or missing_buffers or unexpected_buffers: + raise RuntimeError( + "Backbone and Metis module state do not match: " + f"missing_parameters={missing[:5]}, " + f"unexpected_parameters={unexpected[:5]}, " + f"missing_buffers={missing_buffers[:5]}, " + f"unexpected_buffers={unexpected_buffers[:5]}" + ) + + for name, target_parameter in target_parameters.items(): + source_parameter = source_parameters[name] + target_is_zero3 = _is_zero3_parameter(target_parameter) + source_is_zero3 = _is_zero3_parameter(source_parameter) + if target_is_zero3 != source_is_zero3: + raise RuntimeError( + f"ZeRO-3 state differs for parameter {name}: " + f"target={target_is_zero3}, source={source_is_zero3}" + ) + + if target_is_zero3: + if int(target_parameter.ds_numel) != int(source_parameter.ds_numel): + raise RuntimeError( + f"Shape mismatch for partitioned parameter {name}: " + f"target_numel={target_parameter.ds_numel}, " + f"source_numel={source_parameter.ds_numel}" + ) + target_shard = target_parameter.ds_tensor + source_shard = source_parameter.ds_tensor + if target_shard.shape != source_shard.shape: + raise RuntimeError( + f"Partition mismatch for parameter {name}: " + f"target={tuple(target_shard.shape)}, source={tuple(source_shard.shape)}" + ) + target_shard.copy_( + source_shard.to(device=target_shard.device, dtype=target_shard.dtype) + ) + else: + if target_parameter.shape != source_parameter.shape: + raise RuntimeError( + f"Shape mismatch for parameter {name}: " + f"target={tuple(target_parameter.shape)}, " + f"source={tuple(source_parameter.shape)}" + ) + target_parameter.copy_( + source_parameter.to( + device=target_parameter.device, + dtype=target_parameter.dtype, + ) + ) + + for name, target_buffer in target_buffers.items(): + source_buffer = source_buffers[name] + if target_buffer.shape != source_buffer.shape: + raise RuntimeError( + f"Shape mismatch for buffer {name}: " + f"target={tuple(target_buffer.shape)}, source={tuple(source_buffer.shape)}" + ) + target_buffer.copy_( + source_buffer.to(device=target_buffer.device, dtype=target_buffer.dtype) + ) + + def _fit_rows(src: torch.Tensor, target_rows: int) -> torch.Tensor: """Pad or tile source rows to match the target dimension. @@ -77,32 +244,57 @@ def _init_learned_query_from_backbone(model: MetisForCausalLM) -> None: skipped += 1 continue - src_weight = _extract_backbone_query_rows(src_q_proj.weight, num_heads, head_dim) - if src_weight is None or src_weight.shape != query_proj.weight.shape: - skipped += 1 - continue - - query_proj.weight.copy_(src_weight.to(device=query_proj.weight.device, dtype=query_proj.weight.dtype)) + gathered_parameters = list(src_q_proj.parameters()) + list(query_proj.parameters()) + if query_norm is not None: + gathered_parameters.extend(query_norm.parameters()) + if src_q_norm is not None: + gathered_parameters.extend(src_q_norm.parameters()) + + with _gather_parameters_for_update(gathered_parameters) as should_update: + src_weight = _extract_backbone_query_rows( + src_q_proj.weight, num_heads, head_dim + ) + if src_weight is None or src_weight.shape != query_proj.weight.shape: + skipped += 1 + continue - if query_proj.bias is not None: - src_bias = getattr(src_q_proj, "bias", None) - if src_bias is not None: - src_bias = _extract_backbone_query_rows(src_bias, num_heads, head_dim) - if src_bias is not None and src_bias.shape == query_proj.bias.shape: - query_proj.bias.copy_(src_bias.to(device=query_proj.bias.device, dtype=query_proj.bias.dtype)) - else: + if should_update: + query_proj.weight.copy_( + src_weight.to( + device=query_proj.weight.device, + dtype=query_proj.weight.dtype, + ) + ) + + if query_proj.bias is not None: + src_bias = getattr(src_q_proj, "bias", None) + if src_bias is not None: + src_bias = _extract_backbone_query_rows( + src_bias, num_heads, head_dim + ) + if src_bias is not None and src_bias.shape == query_proj.bias.shape: + if should_update: + query_proj.bias.copy_( + src_bias.to( + device=query_proj.bias.device, + dtype=query_proj.bias.dtype, + ) + ) + elif should_update: + query_proj.bias.zero_() + elif should_update: query_proj.bias.zero_() - else: - query_proj.bias.zero_() - if query_norm is not None and src_q_norm is not None: - try: - query_norm.load_state_dict(src_q_norm.state_dict(), strict=True) - except RuntimeError: - skipped += 1 - continue + if query_norm is not None and src_q_norm is not None: + if not _copy_materialized_module_state( + query_norm, + src_q_norm, + should_update=should_update, + ): + skipped += 1 + continue - inited += 1 + inited += 1 print(f"[weight_utils] Learned memory queries initialized from backbone q_proj/q_norm: {inited} layers (skipped {skipped}).") @@ -130,11 +322,6 @@ def _init_hyper_memory_from_backbone(model: MetisForCausalLM) -> None: continue raw_decoder = block.backbone_decoder.raw_decoder - w_k = hm.W_k.weight # (kv_dim, hidden_size) - w_v = hm.W_v.weight # (kv_dim, hidden_size) - k_rows, k_cols = w_k.shape - v_rows, v_cols = w_v.shape - src_k = src_v = None self_attn = getattr(raw_decoder, "self_attn", None) @@ -148,41 +335,74 @@ def _init_hyper_memory_from_backbone(model: MetisForCausalLM) -> None: in_proj_qkv = getattr(linear_attn, "in_proj_qkv", None) if linear_attn is not None else None in_proj_z = getattr(linear_attn, "in_proj_z", None) if linear_attn is not None else None if in_proj_qkv is not None and in_proj_z is not None: - # in_proj_qkv output layout: [Q(key_dim) | K(key_dim) | V(value_dim)] - qkv_out = in_proj_qkv.weight.shape[0] - value_dim = in_proj_z.weight.shape[0] - key_total = qkv_out - value_dim # key_dim * 2 - - qkv_w = in_proj_qkv.weight - src_k_w = qkv_w[:key_total] # Q+K portion → W_k - src_v_w = qkv_w[key_total:] # V portion → W_v - - if src_k_w.shape[1] != k_cols or src_v_w.shape[1] != v_cols: - skipped += 1 - continue - - w_k.zero_() - w_v.zero_() - w_k[: min(src_k_w.shape[0], k_rows)].copy_(src_k_w[: min(src_k_w.shape[0], k_rows)]) - w_v[: min(src_v_w.shape[0], v_rows)].copy_(src_v_w[: min(src_v_w.shape[0], v_rows)]) - inited += 1 + gathered_parameters = ( + list(hm.W_k.parameters()) + + list(hm.W_v.parameters()) + + list(in_proj_qkv.parameters()) + + list(in_proj_z.parameters()) + ) + with _gather_parameters_for_update(gathered_parameters) as should_update: + w_k = hm.W_k.weight + w_v = hm.W_v.weight + k_rows, k_cols = w_k.shape + v_rows, v_cols = w_v.shape + + # in_proj_qkv layout: + # [Q(key_dim) | K(key_dim) | V(value_dim)] + qkv_out = in_proj_qkv.weight.shape[0] + value_dim = in_proj_z.weight.shape[0] + key_total = qkv_out - value_dim + qkv_w = in_proj_qkv.weight + src_k_w = qkv_w[:key_total] + src_v_w = qkv_w[key_total:] + + if src_k_w.shape[1] != k_cols or src_v_w.shape[1] != v_cols: + skipped += 1 + continue + + if should_update: + w_k.zero_() + w_v.zero_() + w_k[: min(src_k_w.shape[0], k_rows)].copy_( + src_k_w[: min(src_k_w.shape[0], k_rows)] + ) + w_v[: min(src_v_w.shape[0], v_rows)].copy_( + src_v_w[: min(src_v_w.shape[0], v_rows)] + ) + inited += 1 continue skipped += 1 continue - k_w = src_k.weight # (k_out, hidden_size) - v_w = src_v.weight # (v_out, hidden_size) - - if k_w.shape[1] != k_cols or v_w.shape[1] != v_cols: - skipped += 1 - continue + gathered_parameters = ( + list(hm.W_k.parameters()) + + list(hm.W_v.parameters()) + + list(src_k.parameters()) + + list(src_v.parameters()) + ) + with _gather_parameters_for_update(gathered_parameters) as should_update: + w_k = hm.W_k.weight + w_v = hm.W_v.weight + k_rows, k_cols = w_k.shape + v_rows, v_cols = w_v.shape + k_w = src_k.weight + v_w = src_v.weight + + if k_w.shape[1] != k_cols or v_w.shape[1] != v_cols: + skipped += 1 + continue - w_k.zero_() - w_v.zero_() - # Use _fit_rows to handle GQA→MHA dimension mismatch. - w_k[: min(k_w.shape[0], k_rows)].copy_(_fit_rows(k_w, k_rows)[: min(k_w.shape[0], k_rows)]) - w_v[: min(v_w.shape[0], v_rows)].copy_(_fit_rows(v_w, v_rows)[: min(v_w.shape[0], v_rows)]) - inited += 1 + if should_update: + w_k.zero_() + w_v.zero_() + # Use _fit_rows to handle GQA→MHA dimension mismatch. + w_k[: min(k_w.shape[0], k_rows)].copy_( + _fit_rows(k_w, k_rows)[: min(k_w.shape[0], k_rows)] + ) + w_v[: min(v_w.shape[0], v_rows)].copy_( + _fit_rows(v_w, v_rows)[: min(v_w.shape[0], v_rows)] + ) + inited += 1 print(f"[weight_utils] Hyper-memory initialized from backbone projections: {inited} layers (skipped {skipped}).") @@ -269,13 +489,18 @@ def load_metis_from_backbone( }, ) - model = MetisForCausalLM(config) + zero3_enabled = is_deepspeed_zero3_enabled() + if zero3_enabled: + print("[weight_utils] ZeRO-3 detected: constructing Metis directly as parameter shards.") + model = MetisForCausalLM._from_config(config, dtype=dtype) + else: + model = MetisForCausalLM(config) print(f"[weight_utils] Loading backbone weights from {backbone_path} …") backbone = AutoModelForCausalLM.from_pretrained(backbone_path, dtype=dtype) mb = model.model.metis_backbone - mb.model.load_state_dict(backbone.model.state_dict()) - mb.lm_head.load_state_dict(backbone.lm_head.state_dict()) + _copy_module_state(mb.model, backbone.model) + _copy_module_state(mb.lm_head, backbone.lm_head) _init_learned_query_from_backbone(model) _init_hyper_memory_from_backbone(model) @@ -284,7 +509,8 @@ def load_metis_from_backbone( torch.cuda.empty_cache() print("[weight_utils] Backbone weights loaded.") - model = model.to(device=device, dtype=dtype) + if not zero3_enabled: + model = model.to(device=device, dtype=dtype) tokenizer = AutoTokenizer.from_pretrained(backbone_path, trust_remote_code=True) if tokenizer.pad_token is None: diff --git a/scripts/run_train.sh b/scripts/run_train.sh index 6652be4..4560bfe 100755 --- a/scripts/run_train.sh +++ b/scripts/run_train.sh @@ -28,6 +28,7 @@ WARMUP_STEPS=${WARMUP_STEPS:-200} MAX_SEQ_LENGTH=${MAX_SEQ_LENGTH:-4096} MAX_TOTAL_TOKENS=${MAX_TOTAL_TOKENS:-1024} SAVE_STEPS=${SAVE_STEPS:-2000} +SAVE_TOTAL_LIMIT=${SAVE_TOTAL_LIMIT:-2} # Checkpoint format: "delta" = only trainable Metis weights (default), ~100x # smaller since the frozen backbone is not duplicated per checkpoint; @@ -163,6 +164,7 @@ TRAIN_ARGS=( --lora_r 0 --log_steps 5 --save_steps "${SAVE_STEPS}" + --save_total_limit "${SAVE_TOTAL_LIMIT}" --checkpoint_save_mode "${CHECKPOINT_SAVE_MODE}" --eval_steps "${EVAL_STEPS}" --gen_eval_steps "${GEN_EVAL_STEPS}" diff --git a/scripts/train.sh b/scripts/train.sh index 61bf8c8..71a9393 100755 --- a/scripts/train.sh +++ b/scripts/train.sh @@ -28,6 +28,7 @@ # WEIGHT_DECAY=0.01 MAX_GRAD_NORM=1.0 # WARMUP_STEPS=200 NUM_EPOCHS=2 # MAX_STEPS=0 SAVE_STEPS=2000 +# SAVE_TOTAL_LIMIT=2 # CHECKPOINT_SAVE_MODE=delta EVAL_STEPS=1000 # GEN_EVAL_STEPS=0 EVAL_SAMPLES=60 # EVAL_SAMPLES_PER_TASK=20 USE_WANDB=1 @@ -89,6 +90,7 @@ Options: --nproc-per-node N Number of torchrun workers (default: 4). --batch-size N Per-device batch size (default: 2). --grad-accum N Gradient accumulation steps (default: 10). + --save-total-limit N Retain at most N resumable checkpoints (default: 2). --metis-block-type TYPE Metis block implementation. --metis-hyper-memory-type TYPE Hyper-memory implementation. @@ -147,6 +149,8 @@ while [ "$#" -gt 0 ]; do require_value "$@"; BATCH_SIZE=$2; shift 2 ;; --grad-accum) require_value "$@"; GRAD_ACCUM=$2; shift 2 ;; + --save-total-limit) + require_value "$@"; SAVE_TOTAL_LIMIT=$2; shift 2 ;; --metis-block-type) require_value "$@"; METIS_BLOCK_TYPE=$2; shift 2 ;; --metis-hyper-memory-type) @@ -200,6 +204,10 @@ for value_name in NPROC_PER_NODE BATCH_SIZE GRAD_ACCUM; do exit 2 fi done +if [ -n "${SAVE_TOTAL_LIMIT:-}" ] && ! [[ "${SAVE_TOTAL_LIMIT}" =~ ^[0-9]+$ ]]; then + echo "Error: SAVE_TOTAL_LIMIT must be a non-negative integer, got '${SAVE_TOTAL_LIMIT}'." >&2 + exit 2 +fi SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) cd "${SCRIPT_DIR}/.." @@ -270,6 +278,7 @@ export TASK4_WEIGHT_END=${TASK4_WEIGHT_END:-0.15} # keep only the ~126M/202M trainable Metis weights instead of a full model # dump — set CHECKPOINT_SAVE_MODE=full for legacy from_pretrained dumps) ── export SAVE_STEPS=${SAVE_STEPS:-2000} +export SAVE_TOTAL_LIMIT=${SAVE_TOTAL_LIMIT:-2} export CHECKPOINT_SAVE_MODE=${CHECKPOINT_SAVE_MODE:-delta} export EVAL_STEPS=${EVAL_STEPS:-1000} export GEN_EVAL_STEPS=${GEN_EVAL_STEPS:-0} @@ -299,7 +308,7 @@ echo " memory=${METIS_BLOCK_TYPE}/${METIS_HYPER_MEMORY_TYPE}/${METIS_LOCAL_MEM echo " deepspeed=${DS_CONFIG:-off}" echo " gpus=${CUDA_VISIBLE_DEVICES} nproc=${NPROC_PER_NODE} port=${MASTER_PORT}" echo " bs=${BATCH_SIZE} accum=${GRAD_ACCUM} effective=$((NPROC_PER_NODE*BATCH_SIZE*GRAD_ACCUM)) lr=${LR}" -echo " epochs=${NUM_EPOCHS} max_steps=${MAX_STEPS} save_steps=${SAVE_STEPS} ckpt_mode=${CHECKPOINT_SAVE_MODE}" +echo " epochs=${NUM_EPOCHS} max_steps=${MAX_STEPS} save_steps=${SAVE_STEPS} keep_ckpts=${SAVE_TOTAL_LIMIT} ckpt_mode=${CHECKPOINT_SAVE_MODE}" echo " eval_steps=${EVAL_STEPS}(loss) gen_eval_steps=${GEN_EVAL_STEPS}(generation; 0=every eval)" echo " output=${OUTPUT_DIR}/${NAME}" echo " wandb_run=${WANDB_RUN_NAME}" diff --git a/train/dataset.py b/train/dataset.py index b890478..0948bc1 100644 --- a/train/dataset.py +++ b/train/dataset.py @@ -525,7 +525,16 @@ def _summary(self) -> None: def _torch_load(path: Path): + # Tokenized caches are read-only during training. mmap keeps tensor + # storages file-backed, so independent torchrun ranks can share the same + # physical page-cache pages instead of each materializing a private copy + # of every shard in host RAM. Older PyTorch versions (or legacy, + # non-zip torch.save files) do not support mmap; retain a safe fallback. try: - return torch.load(path, map_location="cpu", weights_only=False) + return torch.load(path, map_location="cpu", weights_only=False, mmap=True) except TypeError: return torch.load(path, map_location="cpu") + except RuntimeError as exc: + if "mmap" not in str(exc).lower(): + raise + return torch.load(path, map_location="cpu", weights_only=False) diff --git a/train/run_train.py b/train/run_train.py index 44d9ba2..a100ed8 100644 --- a/train/run_train.py +++ b/train/run_train.py @@ -202,6 +202,50 @@ def _enable_zero3_memory_efficient_linear() -> None: logger.info("Enabled DeepSpeed ZeRO-3 memory_efficient_linear wrapper") +def _build_training_arguments(args) -> TrainingArguments: + """Build Trainer arguments before model construction. + + This ordering is required for Transformers to register + ``HfTrainerDeepSpeedConfig`` before any ``from_pretrained`` call. With + ZeRO-3, model constructors can then partition parameters immediately + instead of first materializing one full model per rank. + """ + return TrainingArguments( + output_dir=args.output_dir, + num_train_epochs=args.num_epochs, + # max_steps > 0 takes precedence over num_train_epochs (HF semantics); + # 0 or negative means "train by epochs". + max_steps=args.max_steps if args.max_steps > 0 else -1, + per_device_train_batch_size=args.batch_size, + gradient_accumulation_steps=args.gradient_accumulation_steps, + learning_rate=args.learning_rate, + weight_decay=args.weight_decay, + warmup_steps=args.warmup_steps, + max_grad_norm=args.max_grad_norm, + adam_beta1=args.adam_beta1, + adam_beta2=args.adam_beta2, + adam_epsilon=args.adam_epsilon, + lr_scheduler_type="constant_with_warmup", + logging_steps=args.log_steps, + eval_strategy="no", # eval is handled via the callback + save_strategy="steps", + save_steps=args.save_steps, + save_total_limit=args.save_total_limit if args.save_total_limit > 0 else None, + save_only_model=False, + bf16=(args.dtype == "bfloat16"), + fp16=(args.dtype == "float16"), + report_to=["wandb"] if args.wandb else ["none"], + run_name=args.wandb_run_name, + remove_unused_columns=False, + dataloader_num_workers=4, + dataloader_pin_memory=True, + ddp_backend="nccl", + ddp_find_unused_parameters=False, + deepspeed=args.deepspeed, + seed=args.seed, + ) + + # ── Task-weighted length sampler ─────────────────────────────────── class TaskWeightedLengthDistributedSampler(Sampler): @@ -532,8 +576,24 @@ def train(args): else: device = torch.device(args.device) + # This must stay alive for the whole load: Transformers keeps only a + # weak reference to the DeepSpeed config consulted by model constructors. + training_args = _build_training_arguments(args) + if args.deepspeed: _enable_zero3_memory_efficient_linear() + from metis.weight_utils import is_deepspeed_zero3_enabled + + if is_deepspeed_zero3_enabled(): + logger.info( + "DeepSpeed ZeRO-3 registered before model load; " + "Metis and backbone parameters will be constructed as shards." + ) + else: + logger.warning( + "DeepSpeed is enabled but stage is not ZeRO-3; " + "model loading will materialize full parameters on each rank." + ) # ── Model ───────────────────────────────────────────────── checkpoint_to_load = args.init_from_checkpoint or args.resume_from_checkpoint @@ -586,6 +646,17 @@ def train(args): qk_kernel_type=args.qk_kernel_type, metis_reweight_gamma=args.metis_reweight_gamma, ) + + if is_main and device.type == "cuda": + gib = 1024**3 + logger.info( + "Model load CUDA memory (rank 0): allocated=%.2f GiB " + "reserved=%.2f GiB peak_allocated=%.2f GiB", + torch.cuda.memory_allocated(device) / gib, + torch.cuda.memory_reserved(device) / gib, + torch.cuda.max_memory_allocated(device) / gib, + ) + if is_main: logger.info( f"memory_configs: " @@ -678,44 +749,6 @@ def train(args): if args.wandb_run_name: os.environ["WANDB_NAME"] = args.wandb_run_name - # ── TrainingArguments ───────────────────────────────────── - train_sampling_kwargs = {} - - training_args = TrainingArguments( - output_dir=args.output_dir, - num_train_epochs=args.num_epochs, - # max_steps > 0 takes precedence over num_train_epochs (HF semantics); - # 0 or negative means "train by epochs". - max_steps=args.max_steps if args.max_steps > 0 else -1, - per_device_train_batch_size=args.batch_size, - gradient_accumulation_steps=args.gradient_accumulation_steps, - learning_rate=args.learning_rate, - weight_decay=args.weight_decay, - warmup_steps=args.warmup_steps, - max_grad_norm=args.max_grad_norm, - adam_beta1=args.adam_beta1, - adam_beta2=args.adam_beta2, - adam_epsilon=args.adam_epsilon, - lr_scheduler_type="constant_with_warmup", - logging_steps=args.log_steps, - eval_strategy="no", # eval is handled via the callback - save_strategy="steps", - save_steps=args.save_steps, - save_only_model=False, - bf16=(args.dtype == "bfloat16"), - fp16=(args.dtype == "float16"), - report_to=["wandb"] if args.wandb else ["none"], - run_name=args.wandb_run_name, - remove_unused_columns=False, - dataloader_num_workers=4, - dataloader_pin_memory=True, - ddp_backend="nccl", - ddp_find_unused_parameters=False, - deepspeed=args.deepspeed, - seed=args.seed, - **train_sampling_kwargs, - ) - # ── Build eval samples (from a small held-out subset) ───── eval_samples = None eval_callback = None @@ -893,16 +926,20 @@ def _float_env(name: str, default: float) -> float: step0_dir = os.path.join(args.output_dir, "checkpoint-0") if args.resume_from_checkpoint: logger.info("Skipping step-0 checkpoint while resuming from %s", args.resume_from_checkpoint) + elif args.deepspeed: + # At this point ZeRO-3 parameters already exist as shards, but Trainer + # has not built its DeepSpeed engine yet. Calling model.state_dict() + # here would either save placeholders or gather the complete 27B model + # on rank 0. Normal step/final saves run after engine initialization. + logger.info( + "Skipping checkpoint-0 before DeepSpeed engine initialization; " + "the first configured step checkpoint will be fully resumable." + ) elif trainer.is_world_process_zero(): os.makedirs(step0_dir, exist_ok=True) - if args.deepspeed: - # DeepSpeed engine is only built inside trainer.train(); pre-train - # checkpoint-0 should save from the still-unwrapped module. - trainer._save(step0_dir, state_dict=trainer.model.state_dict()) - else: - # trainer._save respects checkpoint_save_mode (delta/full) and also - # saves the tokenizer. - trainer._save(step0_dir) + # trainer._save respects checkpoint_save_mode (delta/full) and also + # saves the tokenizer. + trainer._save(step0_dir) logger.info(f"Step-0 checkpoint saved → {step0_dir}") # ── Log task weights per epoch ──────────────────────────── @@ -1003,6 +1040,12 @@ def parse_args(): g.add_argument("--log_file", default="train.log") g.add_argument("--log_steps", type=int, default=5) g.add_argument("--save_steps", type=int, default=500) + g.add_argument( + "--save_total_limit", + type=int, + default=2, + help="Maximum number of resumable checkpoints to retain; 0 keeps all.", + ) g.add_argument("--eval_steps", type=int, default=0, help="Frequency of eval (0 = off). Cheap loss eval runs every this many steps.") g.add_argument("--gen_eval_steps", type=int, default=0, diff --git a/train/train_utils.py b/train/train_utils.py index 062eaeb..7a2aa94 100644 --- a/train/train_utils.py +++ b/train/train_utils.py @@ -96,8 +96,13 @@ def freeze_backbone(model: nn.Module) -> int: def count_params(model: nn.Module) -> tuple[int, int]: - total = sum(p.numel() for p in model.parameters()) - trainable = sum(p.numel() for p in model.parameters() if p.requires_grad) + def logical_numel(parameter: nn.Parameter) -> int: + # ZeRO-3 parameters are placeholders locally; ds_numel records the + # unpartitioned logical size used for meaningful parameter counts. + return int(getattr(parameter, "ds_numel", parameter.numel())) + + total = sum(logical_numel(p) for p in model.parameters()) + trainable = sum(logical_numel(p) for p in model.parameters() if p.requires_grad) return total, trainable @@ -127,8 +132,9 @@ def dump_sampler_tree(dl) -> None: # ═══════════════════════════════════════════════════════════════════ _SNAPSHOT_EXCLUDE_DIRS = { - "experiments", "wandb", "__pycache__", ".git", - "pretrained_models", "data", "checkpoint", "log", + "experiments", "wandb", "__pycache__", ".git", ".pytest_cache", + "pretrained_models", "data", "checkpoint", "checkpoints", + "code_snapshot", "log", "logs", "output", "outputs", } def save_code_snapshot(output_dir: str, project_root: str | None = None) -> None: @@ -136,13 +142,26 @@ def save_code_snapshot(output_dir: str, project_root: str | None = None) -> None if project_root is None: project_root = os.path.dirname(os.path.abspath(__file__)) + project_root = os.path.realpath(project_root) + output_dir = os.path.realpath(output_dir) snap_dir = os.path.join(output_dir, "code_snapshot") os.makedirs(snap_dir, exist_ok=True) + def is_in_output(path: str) -> bool: + try: + return os.path.commonpath((os.path.realpath(path), output_dir)) == output_dir + except ValueError: + return False + copied = 0 for dirpath, dirnames, filenames in os.walk(project_root): - dirnames[:] = [d for d in dirnames - if d not in _SNAPSHOT_EXCLUDE_DIRS and not d.startswith(".")] + dirnames[:] = [ + dirname + for dirname in dirnames + if dirname not in _SNAPSHOT_EXCLUDE_DIRS + and not dirname.startswith(".") + and not is_in_output(os.path.join(dirpath, dirname)) + ] for fname in filenames: if not fname.endswith((".py", ".sh")): continue diff --git a/train/trainer.py b/train/trainer.py index 2d05b09..1614fef 100644 --- a/train/trainer.py +++ b/train/trainer.py @@ -259,12 +259,50 @@ def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys=None) loss = self.compute_loss(model, inputs) return (loss, None, None) + def save_model(self, output_dir=None, _internal_call=False): + """Avoid full-model ZeRO-3 consolidation for delta checkpoints. + + HuggingFace's default DeepSpeed path calls + ``_zero3_consolidated_16bit_state_dict`` before ``_save``. That would + gather the complete frozen 27B model even though a Metis delta needs + only trainable memory parameters. All ranks participate here while + tensors are gathered one at a time. + """ + if not ( + self._checkpoint_save_mode == "delta" + and getattr(self, "is_deepspeed_enabled", False) + ): + return super().save_model(output_dir, _internal_call=_internal_call) + + output_dir = output_dir or self.args.output_dir + unwrapped = self.model.module if hasattr(self.model, "module") else self.model + save_metis_delta_checkpoint( + unwrapped, + output_dir, + tokenizer=self._metis_tokenizer, + base_model_path=self._base_model_path, + state_dict=None, + is_main_process=self.args.should_save, + ) + if self.args.should_save: + self._save_training_logs(output_dir) + if self.args.push_to_hub and not _internal_call: + self.push_to_hub( + commit_message="Model save", + revision=self.args.hub_revision, + ) + def _save(self, output_dir=None, state_dict=None): output_dir = output_dir or self.args.output_dir os.makedirs(output_dir, exist_ok=True) unwrapped = self.model.module if hasattr(self.model, "module") else self.model if state_dict is None and getattr(self, "is_deepspeed_enabled", False): + if self._checkpoint_save_mode == "delta": + raise RuntimeError( + "DeepSpeed delta checkpoints must use save_model() on all ranks " + "to avoid full-model ZeRO-3 consolidation." + ) state_dict = self.accelerator.get_state_dict(self.model_wrapped) if self.args.should_save: From 98c41a02d7dc168c506101fd729a670f77e26daf Mon Sep 17 00:00:00 2001 From: Tong Shen <1833882828@qq.com> Date: Fri, 31 Jul 2026 16:50:40 +0800 Subject: [PATCH 2/2] fix: bug of qwen3.5 wrapper; adjust environment --- README.md | 1 + README_zh.md | 1 + environment.yml | 23 ++++++++--------- metis/backbone_wrappers/Qwen3_5_wrapper.py | 29 +++++++++++++++++++--- 4 files changed, 38 insertions(+), 16 deletions(-) diff --git a/README.md b/README.md index c3661c1..a9b842e 100644 --- a/README.md +++ b/README.md @@ -55,6 +55,7 @@ After each memory step, Metis selects informative hidden states and updates the ```bash conda env create -f environment.yml conda activate metis +python -m pip install --no-build-isolation causal-conv1d==1.6.0 flash-attn==2.7.3 ``` ### 2. Run inference diff --git a/README_zh.md b/README_zh.md index eb7bc16..a837e74 100644 --- a/README_zh.md +++ b/README_zh.md @@ -61,6 +61,7 @@ Metis 仍是一个早期研究系统，并非对外部记忆机制的完整替 ```bash conda env create -f environment.yml conda activate metis +python -m pip install --no-build-isolation causal-conv1d==1.6.0 flash-attn==2.7.3 ``` 仓库中的环境文件是研究环境锁定文件，基于 Python 3.10、PyTorch 2.4.1 + CUDA 11.8、Transformers 5.4.0、DeepSpeed 和 Flash Linear Attention。CUDA 扩展可能需要根据具体硬件与平台进行调整。 diff --git a/environment.yml b/environment.yml index 3ea1317..885b29c 100644 --- a/environment.yml +++ b/environment.yml @@ -1,17 +1,15 @@ name: metis channels: - - conda-forge - nvidia + - conda-forge dependencies: - _openmp_mutex=4.5 - bzip2=1.0.8 - ca-certificates=2026.6.17 - - cuda-cccl=12.9.27 - - cuda-cccl_linux-64=12.9.27 - - cuda-cudart=11.8.89 - - cuda-cudart-dev=11.8.89 - - cuda-nvcc=11.8.89 - - cuda-version=12.9 + - nvidia::cuda-cccl=11.8.89 + - nvidia::cuda-cudart=11.8.89 + - nvidia::cuda-cudart-dev=11.8.89 + - nvidia::cuda-nvcc=11.8.89 - ld_impl_linux-64=2.45.1 - libexpat=2.8.1 - libffi=3.5.2 @@ -36,6 +34,7 @@ dependencies: - zstd=1.5.7 - pip: - --extra-index-url=https://download.pytorch.org/whl/cu118 + # CUDA extensions must be built after torch is installed; see the README. - accelerate==1.14.0 - aiohappyeyeballs==2.6.2 - aiohttp==3.14.1 @@ -45,7 +44,6 @@ dependencies: - anyio==4.14.0 - async-timeout==5.0.1 - attrs==26.1.0 - - causal-conv1d==1.6.2.post1 - certifi==2026.6.17 - charset-normalizer==3.4.7 - click==8.4.1 @@ -57,7 +55,6 @@ dependencies: - einops==0.8.2 - exceptiongroup==1.3.1 - filelock==3.29.0 - - flash-attn==2.6.3 - flash-linear-attention==0.2.2 - fonttools==4.63.0 - frozenlist==1.8.0 @@ -96,7 +93,7 @@ dependencies: - nvidia-cusolver-cu11==11.4.1.48 - nvidia-cusparse-cu11==11.7.5.86 - nvidia-ml-py==13.595.45 - - nvidia-nccl-cu11==2.20.5 + - nvidia-nccl-cu11==2.21.5 - nvidia-nvtx-cu11==11.8.86 - nvitop==1.7.0 - pandas==2.3.3 @@ -124,12 +121,12 @@ dependencies: - six==1.17.0 - smmap==5.0.3 - socksio==1.0.0 - - sympy==1.14.0 + - sympy==1.13.1 - tokenizers==0.22.2 - - torch==2.4.1+cu118 + - torch==2.5.1+cu118 - tqdm==4.68.3 - transformers==5.4.0 - - triton==3.0.0 + - triton==3.1.0 - typer==0.25.1 - typing-extensions==4.15.0 - typing-inspection==0.4.2 diff --git a/metis/backbone_wrappers/Qwen3_5_wrapper.py b/metis/backbone_wrappers/Qwen3_5_wrapper.py index 95cddf0..5aa5534 100644 --- a/metis/backbone_wrappers/Qwen3_5_wrapper.py +++ b/metis/backbone_wrappers/Qwen3_5_wrapper.py @@ -23,6 +23,26 @@ from transformers.modeling_outputs import BaseModelOutputWithPast +def _zero3_safe_checkpoint(function, hidden_states): + """Checkpoint frozen-backbone computation safely under ZeRO-3. + + Reentrant checkpoint avoids the PyTorch 2.5 non-reentrant saved-tensor + metadata conflict with ZeRO-3 parameter re-partitioning. + + If the input does not require gradients, the enclosed backbone operation + is fully frozen, so checkpointing it provides no backward-memory benefit. + """ + if not torch.is_grad_enabled() or not hidden_states.requires_grad: + return function(hidden_states) + + return _ckpt( + function, + hidden_states, + use_reentrant=True, + preserve_rng_state=True, + ) + + class Qwen3_5DecoderLayerForMetis(DecoderLayerWrapperForMetis): def __init__(self, config, raw_decoder): super().__init__(config, raw_decoder) @@ -137,7 +157,10 @@ def _attn(hidden_states): attn_output = self_attn.o_proj(attn_output) return memory_for_query, attn_output, residual - memory_for_query, attn_output, residual = _ckpt(_attn, hidden_states, use_reentrant=False) + memory_for_query, attn_output, residual = _zero3_safe_checkpoint( + _attn, + hidden_states, + ) return memory_for_query, attn_output, {'residual': residual}, self_attn.o_proj @@ -183,7 +206,7 @@ def _mix(hidden_states): ) return residual + mixed - hidden_states = _ckpt(_mix, hidden_states, use_reentrant=False) + hidden_states = _zero3_safe_checkpoint(_mix, hidden_states) elif self.raw_decoder.layer_type == "full_attention": residual = cache_dict['residual'] @@ -194,7 +217,7 @@ def _mlp(hidden_states): normed = self.raw_decoder.post_attention_layernorm(hidden_states) return residual + self.raw_decoder.mlp(normed) - hidden_states = _ckpt(_mlp, hidden_states, use_reentrant=False) + hidden_states = _zero3_safe_checkpoint(_mlp, hidden_states) return hidden_states