From 30629cfa347f6df3005a78ac06a3b392128e9949 Mon Sep 17 00:00:00 2001 From: sunny-infra Date: Sun, 28 Jun 2026 09:09:56 +0000 Subject: [PATCH 1/4] Add Ascend NPU support while preserving GPU compatibility Adapt the training and evaluation framework to run on Ascend NPU (910B2) with no regression on CUDA GPUs. All device selection is runtime-detected via try/except torch_npu with fallback to torch.cuda. Key changes: - distributed: init_dist uses hccl backend + torch.npu on NPU, nccl + torch.cuda otherwise. - utils: seed_all handles NPU RNG; cuda_prefetcher picks npu/cuda Stream. - trainer: FSDP device_mesh selects "npu"/"cuda" device type; RNG state save/load in ckpt_manager supports both backends and tolerates legacy checkpoints; add per-epoch checkpoint saving. - modeling/dspark: add eager attention path (4D additive mask) for NPU since flex_attention is CUDA-only; TRAIN_ATTN_IMPLEMENTATION is chosen at runtime (eager on NPU, flex_attention on GPU). - train.py / eval.py / prepare_target_cache.py: spawn nprocs honors ASCEND_RT_VISIBLE_DEVICES / CUDA_VISIBLE_DEVICES; add --tasks arg to eval.py for selective dataset evaluation. - confidence_head: NPU hccl does not support float64 all_reduce, so cast to float32 for the reduction on NPU only; GPU keeps native float64. Verified on NPU 910B2: target-cache generation, 3-epoch dspark training on Qwen3-8B, and gsm8k speculative-decoding eval all run end-to-end. --- deepspec/data/cuda_prefetcher.py | 23 ++++++++--- deepspec/eval/dspark/confidence_head.py | 19 ++++++++- deepspec/modeling/dspark/common.py | 45 ++++++++++++++++++---- deepspec/modeling/dspark/qwen3/config.py | 14 ++++++- deepspec/modeling/dspark/qwen3/modeling.py | 1 + deepspec/trainer/base_trainer.py | 26 +++++++++++-- deepspec/trainer/ckpt_manager.py | 32 ++++++++++++++- deepspec/utils/__init__.py | 9 ++++- deepspec/utils/distributed.py | 40 ++++++++++++++++--- eval.py | 45 ++++++++++++++++++++-- scripts/data/prepare_target_cache.py | 33 +++++++++++++--- train.py | 22 +++++++++-- 12 files changed, 272 insertions(+), 37 deletions(-) diff --git a/deepspec/data/cuda_prefetcher.py b/deepspec/data/cuda_prefetcher.py index 69eac90..2b95f61 100644 --- a/deepspec/data/cuda_prefetcher.py +++ b/deepspec/data/cuda_prefetcher.py @@ -3,9 +3,21 @@ import torch +def _is_npu_available(): + try: + import torch_npu # noqa: F401 + return torch.npu.is_available() + except Exception: + return False + + +def _stream_module(): + return torch.npu if _is_npu_available() else torch.cuda + + def move_batch_to_device(batch, device): moved = {key: value.to(device, non_blocking=True) for key, value in batch.items()} - # Embedding lookup requires int64; cast on GPU to avoid bloating CPU-to-GPU transfer. + # Embedding lookup requires int64; cast on device to avoid bloating host-to-device transfer. if moved["input_ids"].dtype != torch.long: moved["input_ids"] = moved["input_ids"].to(torch.long) return moved @@ -14,7 +26,7 @@ def move_batch_to_device(batch, device): class CUDAPrefetcher: """Overlaps DataLoader iteration and H2D transfer with compute. - Uses a background thread and a dedicated CUDA stream so that both + Uses a background thread and a dedicated device stream so that both next(dataloader) and the H2D copy for the next batch run concurrently with the forward/backward of the current batch. """ @@ -22,7 +34,8 @@ class CUDAPrefetcher: def __init__(self, dataloader, device): self.dataloader = dataloader self.device = device - self.stream = torch.cuda.Stream(device=device) + self._sm = _stream_module() + self.stream = self._sm.Stream(device=device) def __iter__(self): self._iter = iter(self.dataloader) @@ -40,7 +53,7 @@ def _fetch_and_transfer(self): except StopIteration: self._done = True return - with torch.cuda.stream(self.stream): + with self._sm.stream(self.stream): self._gpu_batch = move_batch_to_device(cpu_batch, self.device) def __next__(self): @@ -53,7 +66,7 @@ def __next__(self): raise StopIteration # Make the compute stream wait until the side-stream H2D is complete. - current = torch.cuda.current_stream(self.device) + current = self._sm.current_stream(self.device) current.wait_stream(self.stream) batch = self._gpu_batch diff --git a/deepspec/eval/dspark/confidence_head.py b/deepspec/eval/dspark/confidence_head.py index 36b1c32..d906739 100644 --- a/deepspec/eval/dspark/confidence_head.py +++ b/deepspec/eval/dspark/confidence_head.py @@ -92,6 +92,10 @@ def update( self.brier_num[:pos_count].add_(weights * (probs - targets).pow(2)) def all_reduce(self) -> None: + # NPU hccl does not support float64 all_reduce (ERR02007). On NPU, + # cast float64 tensors to float32 for the reduction and restore the + # original dtype afterwards. GPU/nccl supports float64 natively. + need_f32_cast = self._on_npu() for tensor in ( self.coarse_count, self.coarse_pred, @@ -100,7 +104,20 @@ def all_reduce(self) -> None: self.fine_neg, self.brier_num, ): - dist.all_reduce(tensor, op=dist.ReduceOp.SUM) + if need_f32_cast and tensor.dtype is torch.float64: + reduced = tensor.to(torch.float32) + dist.all_reduce(reduced, op=dist.ReduceOp.SUM) + tensor.copy_(reduced.to(torch.float64)) + else: + dist.all_reduce(tensor, op=dist.ReduceOp.SUM) + + @staticmethod + def _on_npu() -> bool: + try: + import torch_npu # noqa: F401 + return torch.npu.is_available() + except Exception: + return False @staticmethod def _auroc_from_hist(pos_hist: torch.Tensor, neg_hist: torch.Tensor) -> float: diff --git a/deepspec/modeling/dspark/common.py b/deepspec/modeling/dspark/common.py index 76d908a..9943dfb 100644 --- a/deepspec/modeling/dspark/common.py +++ b/deepspec/modeling/dspark/common.py @@ -82,6 +82,7 @@ def create_dspark_attention_mask( seq_len: int, block_size: int, device: torch.device, + use_block_mask: bool = True, ): def dspark_mask_mod(b, h, q_idx, kv_idx): del h @@ -96,14 +97,42 @@ def dspark_mask_mod(b, h, q_idx, kv_idx): return (mask_context | mask_draft) & is_valid_block bsz, num_blocks = anchor_positions.shape - return create_block_mask( - dspark_mask_mod, - B=bsz, - H=None, - Q_LEN=num_blocks * block_size, - KV_LEN=seq_len + num_blocks * block_size, - device=device, - ) + q_len = num_blocks * block_size + kv_len = seq_len + num_blocks * block_size + if use_block_mask: + return create_block_mask( + dspark_mask_mod, + B=bsz, + H=None, + Q_LEN=q_len, + KV_LEN=kv_len, + device=device, + ) + # Eager / non-flex path: materialize a 4D additive mask (bool -> 0/-inf). + # attn pattern: (mask_context | mask_draft) & is_valid_block + q_idx = torch.arange(q_len, device=device).view(1, 1, q_len, 1) + kv_idx = torch.arange(kv_len, device=device).view(1, 1, 1, kv_len) + q_block_id = q_idx // block_size # (1,1,q_len,1) + # gather anchor_pos per q_block_id: (bsz,1,q_len,1) + anchor_pos_q = torch.gather( + anchor_positions, dim=1, index=q_block_id.view(bsz, -1) + ).view(bsz, 1, q_len, 1) + is_context = kv_idx < seq_len + mask_context = is_context & (kv_idx < anchor_pos_q) + is_draft = kv_idx >= seq_len + kv_block_id = (kv_idx - seq_len) // block_size + mask_draft = is_draft & (q_block_id == kv_block_id) + # is_valid_block per q_block_id -> (bsz,1,q_len,1) + is_valid_block_q = torch.gather( + block_keep_mask, dim=1, index=q_block_id.view(bsz, -1) + ).view(bsz, 1, q_len, 1) + keep = (mask_context | mask_draft) & is_valid_block_q # bool (bsz,1,q_len,kv_len) + # Convert to additive mask: 0 where keep, -inf elsewhere. eager_attention_forward + # does `attn_weights + attention_mask`, so use min dtype-safe value. + neg_inf = torch.finfo(torch.float32).min + additive = torch.zeros(bsz, 1, q_len, kv_len, dtype=torch.float32, device=device) + additive = additive.masked_fill(~keep, neg_inf) + return additive def build_anchor_candidate_mask( diff --git a/deepspec/modeling/dspark/qwen3/config.py b/deepspec/modeling/dspark/qwen3/config.py index ae71fac..ab0142b 100644 --- a/deepspec/modeling/dspark/qwen3/config.py +++ b/deepspec/modeling/dspark/qwen3/config.py @@ -1,9 +1,21 @@ import copy +import torch + from deepspec.modeling.dspark.common import validate_target_layer_ids -TRAIN_ATTN_IMPLEMENTATION = "flex_attention" +def _is_npu_available(): + try: + import torch_npu # noqa: F401 + return torch.npu.is_available() + except Exception: + return False + + +# NPU does not support flex_attention; fall back to eager. GPU keeps the +# original flex_attention path for best performance. +TRAIN_ATTN_IMPLEMENTATION = "eager" if _is_npu_available() else "flex_attention" def build_draft_config( diff --git a/deepspec/modeling/dspark/qwen3/modeling.py b/deepspec/modeling/dspark/qwen3/modeling.py index 2f9c220..b42fac8 100644 --- a/deepspec/modeling/dspark/qwen3/modeling.py +++ b/deepspec/modeling/dspark/qwen3/modeling.py @@ -419,6 +419,7 @@ def forward( seq_len=seq_len, block_size=self.block_size, device=device, + use_block_mask=self.config._attn_implementation == "flex_attention", ) output_hidden = self._forward_backbone( position_ids=full_position_ids, diff --git a/deepspec/trainer/base_trainer.py b/deepspec/trainer/base_trainer.py index 1a03dc8..d9a0f6e 100644 --- a/deepspec/trainer/base_trainer.py +++ b/deepspec/trainer/base_trainer.py @@ -52,6 +52,14 @@ ) +def _is_npu_available(): + try: + import torch_npu # noqa: F401 + return torch.npu.is_available() + except Exception: + return False + + def _build_fsdp_kwargs( *, sharding_strategy_name: str, precision_dtype, world_size: int ) -> dict: @@ -65,9 +73,14 @@ def _build_fsdp_kwargs( sharding_strategy=sharding_strategy, ) if sharding_strategy in _HYBRID_STRATEGIES: - devices_per_node = torch.cuda.device_count() + if _is_npu_available(): + devices_per_node = torch.npu.device_count() + device_type = "npu" + else: + devices_per_node = torch.cuda.device_count() + device_type = "cuda" fsdp_kwargs["device_mesh"] = init_device_mesh( - "cuda", + device_type, (world_size // devices_per_node, devices_per_node), mesh_dim_names=("replicate", "shard"), ) @@ -390,7 +403,14 @@ def train(self): grad_norm=grad_norm.item(), ) - if self.global_step % int(self.args.logging.checkpointing_steps) == 0: + should_checkpoint = ( + self.global_step % int(self.args.logging.checkpointing_steps) == 0 + or ( + self.global_step % self.steps_per_epoch == 0 + and self.global_step < self.max_train_steps + ) + ) + if should_checkpoint: self.save_and_eval_checkpoint() if self.suspend_controller.requested(): diff --git a/deepspec/trainer/ckpt_manager.py b/deepspec/trainer/ckpt_manager.py index 6b99f65..d04a92f 100644 --- a/deepspec/trainer/ckpt_manager.py +++ b/deepspec/trainer/ckpt_manager.py @@ -112,7 +112,7 @@ def load_training_state( assert saved_local_batch_size == int(local_batch_size) torch.set_rng_state(checkpoint["torch_rng"]) - torch.cuda.set_rng_state(checkpoint["torch_cuda_rng"]) + _set_device_rng_state(checkpoint.get("torch_cuda_rng")) np.random.set_state(checkpoint["numpy_rng"]) random.setstate(checkpoint["python_rng"]) @@ -185,6 +185,34 @@ def save_checkpoint( return checkpoint_dir +def _get_device_rng_state(): + """Return RNG state of the active accelerator (NPU if available, else CUDA).""" + try: + import torch_npu # noqa: F401 + if torch.npu.is_available(): + return torch.npu.get_rng_state() + except Exception: + pass + if torch.cuda.is_available(): + return torch.cuda.get_rng_state() + return None + + +def _set_device_rng_state(state): + """Restore RNG state on the active accelerator (NPU if available, else CUDA).""" + if state is None: + return + try: + import torch_npu # noqa: F401 + if torch.npu.is_available(): + torch.npu.set_rng_state(state) + return + except Exception: + pass + if torch.cuda.is_available(): + torch.cuda.set_rng_state(state) + + def _rank_training_state_path(checkpoint_dir: str, global_rank: int) -> str: return os.path.join( checkpoint_dir, @@ -213,7 +241,7 @@ def _serialize_training_state( "world_size": int(world_size), "local_batch_size": int(local_batch_size), "torch_rng": torch.get_rng_state(), - "torch_cuda_rng": torch.cuda.get_rng_state(), + "torch_cuda_rng": _get_device_rng_state(), "numpy_rng": np.random.get_state(), "python_rng": random.getstate(), } diff --git a/deepspec/utils/__init__.py b/deepspec/utils/__init__.py index cdf9aff..c1145d1 100644 --- a/deepspec/utils/__init__.py +++ b/deepspec/utils/__init__.py @@ -21,7 +21,14 @@ def seed_all(seed): torch.manual_seed(seed) random.seed(seed) np.random.seed(seed) - torch.cuda.manual_seed_all(seed) + try: + import torch_npu # noqa: F401 + torch.npu.manual_seed_all(seed) + return + except Exception: + pass + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) def get_git_sha(detail_info=False): import subprocess diff --git a/deepspec/utils/distributed.py b/deepspec/utils/distributed.py index 868bb2d..a9193db 100644 --- a/deepspec/utils/distributed.py +++ b/deepspec/utils/distributed.py @@ -9,18 +9,47 @@ from torch.utils.data import Sampler +def _is_npu_available(): + try: + import torch_npu # noqa: F401 + return torch.npu.is_available() + except Exception: + return False + + +def _device_backend_info(): + """Return (device_module, device_count_fn, set_device_fn, current_device_fn, backend).""" + if _is_npu_available(): + return ( + torch.npu, + torch.npu.device_count, + torch.npu.set_device, + torch.npu.current_device, + "hccl", + ) + return ( + torch.cuda, + torch.cuda.device_count, + torch.cuda.set_device, + torch.cuda.current_device, + "nccl", + ) + + def init_dist(local_rank: int, timeout_minutes: int = 60): - local_world_size = torch.cuda.device_count() + dev_mod, device_count, set_device, _current_device, backend = _device_backend_info() + local_world_size = device_count() node_rank = int(os.environ["RANK"]) node_world_size = int(os.environ["WORLD_SIZE"]) rank = node_rank * local_world_size + local_rank world_size = node_world_size * local_world_size init_method = f"tcp://{os.environ['MASTER_ADDR']}:{os.environ['MASTER_PORT']}" - torch.cuda.set_device(local_rank) - device = torch.device("cuda", local_rank) + set_device(local_rank) + device_type = "npu" if _is_npu_available() else "cuda" + device = torch.device(device_type, local_rank) dist.init_process_group( - backend="nccl", + backend=backend, init_method=init_method, rank=rank, world_size=world_size, @@ -35,7 +64,8 @@ def is_global_main_process(): def is_local_main_process(): - return torch.cuda.current_device() == 0 + _dev_mod, _dc, _sd, current_device, _be = _device_backend_info() + return current_device() == 0 def print_on_global_main(*args, **kwargs): diff --git a/eval.py b/eval.py index a35e7ea..ecd5261 100644 --- a/eval.py +++ b/eval.py @@ -1,6 +1,8 @@ from __future__ import annotations import argparse import json +import os + import torch from transformers import AutoConfig from deepspec.eval.dspark import Gemma4DSparkEvaluator, Qwen3DSparkEvaluator @@ -16,7 +18,7 @@ } TASKS = [ - ("gsm8k", 500), + ("gsm8k", 10), ("math500", 500), ("aime25",30), ("humaneval", 164), @@ -27,6 +29,32 @@ ("arena-hard-v2", 500), ] +def _is_npu_available(): + try: + import torch_npu # noqa: F401 + return torch.npu.is_available() + except Exception: + return False + + +def _visible_device_count(): + """Return the number of devices actually usable by this process. + + Honors ``ASCEND_RT_VISIBLE_DEVICES`` (NPU) and ``CUDA_VISIBLE_DEVICES`` + (CUDA) so that ``torch.multiprocessing.spawn`` does not over-spawn + workers when the process is restricted to a subset of physical devices. + """ + if _is_npu_available(): + vis = os.environ.get("ASCEND_RT_VISIBLE_DEVICES", "").strip() + if vis: + return len([d for d in vis.split(",") if d.strip() != ""]) + return torch.npu.device_count() + vis = os.environ.get("CUDA_VISIBLE_DEVICES", "").strip() + if vis: + return len([d for d in vis.split(",") if d.strip() != ""]) + return torch.cuda.device_count() + + def parse_args(): parser = argparse.ArgumentParser() parser.add_argument("--target_name_or_path", type=str, required=True) @@ -42,8 +70,19 @@ def parse_args(): parser.add_argument("--tensorboard-dir", type=str, default=None) parser.add_argument("--step", type=int, default=None,help=("step for tensorboard logging"),) parser.add_argument("--seed", type=int, default=980406) + parser.add_argument( + "--tasks", + type=str, + default=None, + help=("Comma-separated dataset names to evaluate (e.g. gsm8k). Defaults to all tasks."), + ) args = parser.parse_args() - args.tasks = list(TASKS) + if args.tasks is not None: + selected = {name.strip() for name in args.tasks.split(",")} + args.tasks = [task for task in TASKS if task[0] in selected] + assert args.tasks, f"No matching tasks for --tasks={args.tasks}" + else: + args.tasks = list(TASKS) return args @@ -61,5 +100,5 @@ def main(local_rank: int, args): torch.multiprocessing.spawn( main, args=(args,), - nprocs=torch.cuda.device_count(), + nprocs=_visible_device_count(), ) diff --git a/scripts/data/prepare_target_cache.py b/scripts/data/prepare_target_cache.py index d072582..744d611 100644 --- a/scripts/data/prepare_target_cache.py +++ b/scripts/data/prepare_target_cache.py @@ -46,6 +46,27 @@ torch.set_float32_matmul_precision("high") +def _is_npu_available(): + try: + import torch_npu # noqa: F401 + return torch.npu.is_available() + except Exception: + return False + + +def _device_count(): + if _is_npu_available(): + return torch.npu.device_count() + return torch.cuda.device_count() + + +def _empty_cache(): + if _is_npu_available(): + torch.npu.empty_cache() + elif torch.cuda.is_available(): + torch.cuda.empty_cache() + + @dataclass(frozen=True) class TargetForwardResult: target_hidden_states: torch.Tensor @@ -254,7 +275,7 @@ def main(local_rank: int): target_model = AutoModel.from_pretrained( config.model.target_model_name_or_path, dtype=torch.bfloat16, - attn_implementation="sdpa", + attn_implementation="eager", ).to(device=device).eval() target_hidden_size = _get_target_hidden_size(target_model) train_collator = ConversationCollator( @@ -335,7 +356,7 @@ def main(local_rank: int): finally: writer.close() del target_model - torch.cuda.empty_cache() + _empty_cache() dataset.close() summary = LocalCacheWriteSummary( global_rank=global_rank, @@ -397,6 +418,8 @@ def main(local_rank: int): if __name__ == "__main__": if os.path.exists(".git"): - print(f"git status:", "\n\n".join(get_git_sha(detail_info=True))) - print("git diff:", get_git_diff()) - torch.multiprocessing.spawn(main, nprocs=torch.cuda.device_count()) + try: + print(f"git sha:", get_git_sha()) + except Exception: + pass + torch.multiprocessing.spawn(main, nprocs=_device_count()) diff --git a/train.py b/train.py index 40d9e07..c8ccfcd 100644 --- a/train.py +++ b/train.py @@ -17,6 +17,20 @@ torch.set_float32_matmul_precision("high") +def _is_npu_available(): + try: + import torch_npu # noqa: F401 + return torch.npu.is_available() + except Exception: + return False + + +def _device_count(): + if _is_npu_available(): + return torch.npu.device_count() + return torch.cuda.device_count() + + def parse_args(): parser = argparse.ArgumentParser() parser.add_argument("--config", required=True) @@ -40,6 +54,8 @@ def main(local_rank): if __name__ == "__main__": if os.path.exists(".git"): - print(f"git status:", "\n\n".join(get_git_sha(detail_info=True))) - print("git diff:", get_git_diff()) - torch.multiprocessing.spawn(main, nprocs=torch.cuda.device_count()) + try: + print(f"git sha:", get_git_sha()) + except Exception: + pass + torch.multiprocessing.spawn(main, nprocs=_device_count()) From 53167f57624a7bee582c50d51c2cb7f8668034e9 Mon Sep 17 00:00:00 2001 From: hazel duan Date: Mon, 29 Jun 2026 12:06:37 +0800 Subject: [PATCH 2/4] fix: align Ascend DSpark compatibility --- deepspec/data/cuda_prefetcher.py | 24 ++--- deepspec/eval/dspark/confidence_head.py | 98 +++++++++------------ deepspec/modeling/dspark/common.py | 72 ++++++++------- deepspec/modeling/dspark/gemma4/config.py | 8 +- deepspec/modeling/dspark/gemma4/modeling.py | 1 + deepspec/modeling/dspark/qwen3/config.py | 15 +--- deepspec/modeling/dspark/qwen3/modeling.py | 2 +- deepspec/trainer/base_trainer.py | 29 ++---- deepspec/trainer/ckpt_manager.py | 36 ++------ deepspec/utils/__init__.py | 35 ++++++-- deepspec/utils/device.py | 76 ++++++++++++++++ deepspec/utils/distributed.py | 49 ++++------- deepspec/utils/metrics.py | 6 +- eval.py | 62 +++++-------- requirements.txt | 1 + scripts/data/prepare_target_cache.py | 35 ++------ scripts/train/train.sh | 75 ++++++++-------- train.py | 23 +---- 18 files changed, 304 insertions(+), 343 deletions(-) create mode 100644 deepspec/utils/device.py mode change 100644 => 100755 scripts/train/train.sh diff --git a/deepspec/data/cuda_prefetcher.py b/deepspec/data/cuda_prefetcher.py index 2b95f61..7f682a4 100644 --- a/deepspec/data/cuda_prefetcher.py +++ b/deepspec/data/cuda_prefetcher.py @@ -2,22 +2,12 @@ import torch - -def _is_npu_available(): - try: - import torch_npu # noqa: F401 - return torch.npu.is_available() - except Exception: - return False - - -def _stream_module(): - return torch.npu if _is_npu_available() else torch.cuda +from deepspec.utils import accelerator_module def move_batch_to_device(batch, device): moved = {key: value.to(device, non_blocking=True) for key, value in batch.items()} - # Embedding lookup requires int64; cast on device to avoid bloating host-to-device transfer. + # Embedding lookup requires int64; cast on GPU to avoid bloating CPU-to-GPU transfer. if moved["input_ids"].dtype != torch.long: moved["input_ids"] = moved["input_ids"].to(torch.long) return moved @@ -26,7 +16,7 @@ def move_batch_to_device(batch, device): class CUDAPrefetcher: """Overlaps DataLoader iteration and H2D transfer with compute. - Uses a background thread and a dedicated device stream so that both + Uses a background thread and a dedicated CUDA stream so that both next(dataloader) and the H2D copy for the next batch run concurrently with the forward/backward of the current batch. """ @@ -34,8 +24,8 @@ class CUDAPrefetcher: def __init__(self, dataloader, device): self.dataloader = dataloader self.device = device - self._sm = _stream_module() - self.stream = self._sm.Stream(device=device) + self.accelerator = accelerator_module() + self.stream = self.accelerator.Stream(device=device) def __iter__(self): self._iter = iter(self.dataloader) @@ -53,7 +43,7 @@ def _fetch_and_transfer(self): except StopIteration: self._done = True return - with self._sm.stream(self.stream): + with self.accelerator.stream(self.stream): self._gpu_batch = move_batch_to_device(cpu_batch, self.device) def __next__(self): @@ -66,7 +56,7 @@ def __next__(self): raise StopIteration # Make the compute stream wait until the side-stream H2D is complete. - current = self._sm.current_stream(self.device) + current = self.accelerator.current_stream(self.device) current.wait_stream(self.stream) batch = self._gpu_batch diff --git a/deepspec/eval/dspark/confidence_head.py b/deepspec/eval/dspark/confidence_head.py index d906739..5380a68 100644 --- a/deepspec/eval/dspark/confidence_head.py +++ b/deepspec/eval/dspark/confidence_head.py @@ -27,6 +27,11 @@ def model_display_name(path: str) -> str: return os.path.basename(normalized) or normalized +def confidence_metric_dtype(device: torch.device) -> torch.dtype: + device = torch.device(device) + return torch.float32 if device.type == "npu" else torch.float64 + + class PerPositionConfidenceMetrics: """Per-position ECE + AUROC + Brier for cumprod predictions.""" @@ -41,22 +46,23 @@ def __init__( self.block_size = int(block_size) self.num_coarse_bins = int(num_coarse_bins) self.num_fine_bins = int(num_fine_bins) + self.dtype = confidence_metric_dtype(device) self.coarse_count = torch.zeros( (self.block_size, self.num_coarse_bins), - dtype=torch.float64, + dtype=self.dtype, device=device, ) self.coarse_pred = torch.zeros_like(self.coarse_count) self.coarse_target = torch.zeros_like(self.coarse_count) self.fine_pos = torch.zeros( (self.block_size, self.num_fine_bins), - dtype=torch.float64, + dtype=self.dtype, device=device, ) self.fine_neg = torch.zeros_like(self.fine_pos) self.brier_num = torch.zeros( self.block_size, - dtype=torch.float64, + dtype=self.dtype, device=device, ) @@ -66,12 +72,12 @@ def update( probs: torch.Tensor, targets: torch.Tensor, ) -> None: - probs = probs.reshape(-1).to(torch.float64).clamp(EPS_PROB, 1.0 - EPS_PROB) - targets = targets.reshape(-1).to(torch.float64) + probs = probs.reshape(-1).to(self.dtype).clamp(EPS_PROB, 1.0 - EPS_PROB) + targets = targets.reshape(-1).to(self.dtype) assert probs.shape == targets.shape pos_count = probs.numel() assert pos_count <= self.block_size - weights = torch.ones_like(probs, dtype=torch.float64) + weights = torch.ones_like(probs, dtype=self.dtype) pos_idx = torch.arange(pos_count, device=probs.device) coarse_idx = ( @@ -92,10 +98,6 @@ def update( self.brier_num[:pos_count].add_(weights * (probs - targets).pow(2)) def all_reduce(self) -> None: - # NPU hccl does not support float64 all_reduce (ERR02007). On NPU, - # cast float64 tensors to float32 for the reduction and restore the - # original dtype afterwards. GPU/nccl supports float64 natively. - need_f32_cast = self._on_npu() for tensor in ( self.coarse_count, self.coarse_pred, @@ -104,20 +106,7 @@ def all_reduce(self) -> None: self.fine_neg, self.brier_num, ): - if need_f32_cast and tensor.dtype is torch.float64: - reduced = tensor.to(torch.float32) - dist.all_reduce(reduced, op=dist.ReduceOp.SUM) - tensor.copy_(reduced.to(torch.float64)) - else: - dist.all_reduce(tensor, op=dist.ReduceOp.SUM) - - @staticmethod - def _on_npu() -> bool: - try: - import torch_npu # noqa: F401 - return torch.npu.is_available() - except Exception: - return False + dist.all_reduce(tensor, op=dist.ReduceOp.SUM) @staticmethod def _auroc_from_hist(pos_hist: torch.Tensor, neg_hist: torch.Tensor) -> float: @@ -380,11 +369,11 @@ def observe( step_probs = torch.sigmoid( confidence_logits[:, :effective_length] ).squeeze(0) - cumprod_pred = step_probs.to(torch.float64).cumprod(dim=0) + cumprod_pred = step_probs.to(self.dataset_metrics.dtype).cumprod(dim=0) prefix_label = ( verification.accept_prefix_mask[:, :effective_length] .squeeze(0) - .to(torch.float64) + .to(self.dataset_metrics.dtype) ) self.dataset_metrics.update( probs=cumprod_pred, @@ -585,35 +574,34 @@ def build_table(self) -> str: int(entry["position"]): entry for entry in row.get("per_position") or [] } - table.add_row( - [ - row["dataset"], - draft_name, - row["sample_count"], - row["proposal_count"], - format_float(summary["ece_mean"]), - format_float(summary["auc_mean"]), - format_float(summary["brier_mean"]), - format_float(summary["pred_mean"]), - format_float(summary["target_mean"]), - ] + [ - ( - format_float(per_position[pos]["ece"]) - if pos in per_position - and float(per_position[pos]["total_weight"]) > 0.0 - else "-" - ) - for pos in range(max_position_count) - ] + [ - ( - format_float(per_position[pos]["auc"]) - if pos in per_position - and float(per_position[pos]["total_weight"]) > 0.0 - else "-" - ) - for pos in range(max_position_count) - ] - ) + table_row = [ + row["dataset"], + draft_name, + row["sample_count"], + row["proposal_count"], + format_float(summary["ece_mean"]), + format_float(summary["auc_mean"]), + format_float(summary["brier_mean"]), + format_float(summary["pred_mean"]), + format_float(summary["target_mean"]), + ] + [ + ( + format_float(per_position[pos]["ece"]) + if pos in per_position + and float(per_position[pos]["total_weight"]) > 0.0 + else "-" + ) + for pos in range(max_position_count) + ] + [ + ( + format_float(per_position[pos]["auc"]) + if pos in per_position + and float(per_position[pos]["total_weight"]) > 0.0 + else "-" + ) + for pos in range(max_position_count) + ] + table.add_row(table_row) return table.get_string() def print_results(self) -> None: diff --git a/deepspec/modeling/dspark/common.py b/deepspec/modeling/dspark/common.py index 9943dfb..1a6f6fb 100644 --- a/deepspec/modeling/dspark/common.py +++ b/deepspec/modeling/dspark/common.py @@ -82,8 +82,34 @@ def create_dspark_attention_mask( seq_len: int, block_size: int, device: torch.device, - use_block_mask: bool = True, + attn_implementation: str = "flex_attention", ): + if attn_implementation != "flex_attention": + bsz, num_blocks = anchor_positions.shape + q_len = num_blocks * block_size + kv_len = seq_len + q_len + q_idx = torch.arange(q_len, device=device) + kv_idx = torch.arange(kv_len, device=device) + q_block_ids = (q_idx // block_size).unsqueeze(0).expand(bsz, -1) + anchor_pos = anchor_positions.gather(1, q_block_ids).unsqueeze(-1) + q_block_ids = q_block_ids.unsqueeze(-1) + kv_idx = kv_idx.view(1, 1, kv_len) + + is_context = kv_idx < seq_len + mask_context = is_context & (kv_idx < anchor_pos) + is_draft = kv_idx >= seq_len + kv_block_ids = (kv_idx - seq_len) // block_size + mask_draft = is_draft & (q_block_ids == kv_block_ids) + is_valid_block = block_keep_mask.gather( + 1, + q_block_ids.squeeze(-1), + ).unsqueeze(-1) + dense_mask = (mask_context | mask_draft) & is_valid_block + empty_rows = ~dense_mask.any(dim=-1, keepdim=True) + self_kv_idx = int(seq_len) + q_idx.view(1, -1, 1) + dense_mask = dense_mask | (empty_rows & (kv_idx == self_kv_idx)) + return dense_mask.unsqueeze(1) + def dspark_mask_mod(b, h, q_idx, kv_idx): del h q_block_id = q_idx // block_size @@ -97,42 +123,14 @@ def dspark_mask_mod(b, h, q_idx, kv_idx): return (mask_context | mask_draft) & is_valid_block bsz, num_blocks = anchor_positions.shape - q_len = num_blocks * block_size - kv_len = seq_len + num_blocks * block_size - if use_block_mask: - return create_block_mask( - dspark_mask_mod, - B=bsz, - H=None, - Q_LEN=q_len, - KV_LEN=kv_len, - device=device, - ) - # Eager / non-flex path: materialize a 4D additive mask (bool -> 0/-inf). - # attn pattern: (mask_context | mask_draft) & is_valid_block - q_idx = torch.arange(q_len, device=device).view(1, 1, q_len, 1) - kv_idx = torch.arange(kv_len, device=device).view(1, 1, 1, kv_len) - q_block_id = q_idx // block_size # (1,1,q_len,1) - # gather anchor_pos per q_block_id: (bsz,1,q_len,1) - anchor_pos_q = torch.gather( - anchor_positions, dim=1, index=q_block_id.view(bsz, -1) - ).view(bsz, 1, q_len, 1) - is_context = kv_idx < seq_len - mask_context = is_context & (kv_idx < anchor_pos_q) - is_draft = kv_idx >= seq_len - kv_block_id = (kv_idx - seq_len) // block_size - mask_draft = is_draft & (q_block_id == kv_block_id) - # is_valid_block per q_block_id -> (bsz,1,q_len,1) - is_valid_block_q = torch.gather( - block_keep_mask, dim=1, index=q_block_id.view(bsz, -1) - ).view(bsz, 1, q_len, 1) - keep = (mask_context | mask_draft) & is_valid_block_q # bool (bsz,1,q_len,kv_len) - # Convert to additive mask: 0 where keep, -inf elsewhere. eager_attention_forward - # does `attn_weights + attention_mask`, so use min dtype-safe value. - neg_inf = torch.finfo(torch.float32).min - additive = torch.zeros(bsz, 1, q_len, kv_len, dtype=torch.float32, device=device) - additive = additive.masked_fill(~keep, neg_inf) - return additive + return create_block_mask( + dspark_mask_mod, + B=bsz, + H=None, + Q_LEN=num_blocks * block_size, + KV_LEN=seq_len + num_blocks * block_size, + device=device, + ) def build_anchor_candidate_mask( diff --git a/deepspec/modeling/dspark/gemma4/config.py b/deepspec/modeling/dspark/gemma4/config.py index 3c10f16..a92a8a4 100644 --- a/deepspec/modeling/dspark/gemma4/config.py +++ b/deepspec/modeling/dspark/gemma4/config.py @@ -1,9 +1,10 @@ import copy from deepspec.modeling.dspark.common import validate_target_layer_ids +from deepspec.utils import is_npu_available -TRAIN_ATTN_IMPLEMENTATION = "flex_attention" +TRAIN_ATTN_IMPLEMENTATION = "sdpa" if is_npu_available() else "flex_attention" def get_gemma4_text_config(target_config): @@ -84,6 +85,11 @@ def build_draft_config(target_config, model_args): draft_config.target_text_model_type = str(draft_config.model_type) draft_config.num_target_layers = num_target_layers draft_config.num_hidden_layers = num_draft_layers + # The public Gemma4 DSpark draft block is dense-only. Gemma4 A4B targets are + # MoE, but the draft model can still train as a dense proposal network + # against target hidden states from the MoE model. + draft_config.enable_moe_block = False + draft_config.hidden_size_per_layer_input = 0 draft_config.block_size = int(model_args.block_size) draft_config.tie_word_embeddings = False draft_config.layer_types = layer_types diff --git a/deepspec/modeling/dspark/gemma4/modeling.py b/deepspec/modeling/dspark/gemma4/modeling.py index c339b97..b4fbb8a 100644 --- a/deepspec/modeling/dspark/gemma4/modeling.py +++ b/deepspec/modeling/dspark/gemma4/modeling.py @@ -484,6 +484,7 @@ def forward( seq_len=seq_len, block_size=self.block_size, device=device, + attn_implementation=self.config._attn_implementation, ) output_hidden = self._forward_backbone( position_ids=full_position_ids, diff --git a/deepspec/modeling/dspark/qwen3/config.py b/deepspec/modeling/dspark/qwen3/config.py index ab0142b..03ccf34 100644 --- a/deepspec/modeling/dspark/qwen3/config.py +++ b/deepspec/modeling/dspark/qwen3/config.py @@ -1,21 +1,10 @@ import copy -import torch - from deepspec.modeling.dspark.common import validate_target_layer_ids +from deepspec.utils import is_npu_available -def _is_npu_available(): - try: - import torch_npu # noqa: F401 - return torch.npu.is_available() - except Exception: - return False - - -# NPU does not support flex_attention; fall back to eager. GPU keeps the -# original flex_attention path for best performance. -TRAIN_ATTN_IMPLEMENTATION = "eager" if _is_npu_available() else "flex_attention" +TRAIN_ATTN_IMPLEMENTATION = "sdpa" if is_npu_available() else "flex_attention" def build_draft_config( diff --git a/deepspec/modeling/dspark/qwen3/modeling.py b/deepspec/modeling/dspark/qwen3/modeling.py index b42fac8..3c3dd01 100644 --- a/deepspec/modeling/dspark/qwen3/modeling.py +++ b/deepspec/modeling/dspark/qwen3/modeling.py @@ -419,7 +419,7 @@ def forward( seq_len=seq_len, block_size=self.block_size, device=device, - use_block_mask=self.config._attn_implementation == "flex_attention", + attn_implementation=self.config._attn_implementation, ) output_hidden = self._forward_backbone( position_ids=full_position_ids, diff --git a/deepspec/trainer/base_trainer.py b/deepspec/trainer/base_trainer.py index d9a0f6e..4e6c1a6 100644 --- a/deepspec/trainer/base_trainer.py +++ b/deepspec/trainer/base_trainer.py @@ -15,6 +15,8 @@ from deepspec.utils import ( BF16Optimizer, StatelessResumableDistributedSampler, + device_count, + device_type, ensure_dir, init_dist, is_global_main_process, @@ -52,14 +54,6 @@ ) -def _is_npu_available(): - try: - import torch_npu # noqa: F401 - return torch.npu.is_available() - except Exception: - return False - - def _build_fsdp_kwargs( *, sharding_strategy_name: str, precision_dtype, world_size: int ) -> dict: @@ -73,14 +67,9 @@ def _build_fsdp_kwargs( sharding_strategy=sharding_strategy, ) if sharding_strategy in _HYBRID_STRATEGIES: - if _is_npu_available(): - devices_per_node = torch.npu.device_count() - device_type = "npu" - else: - devices_per_node = torch.cuda.device_count() - device_type = "cuda" + devices_per_node = device_count() fsdp_kwargs["device_mesh"] = init_device_mesh( - device_type, + device_type(), (world_size // devices_per_node, devices_per_node), mesh_dim_names=("replicate", "shard"), ) @@ -296,6 +285,7 @@ def _wrap_with_fsdp(self, model): precision_dtype=self.precision_dtype, world_size=self.world_size, ) + fsdp_kwargs["device_id"] = self.device return FSDP(model, **fsdp_kwargs) def _build_train_dataloader(self, start_offset_samples=0, num_samples=None): @@ -403,14 +393,7 @@ def train(self): grad_norm=grad_norm.item(), ) - should_checkpoint = ( - self.global_step % int(self.args.logging.checkpointing_steps) == 0 - or ( - self.global_step % self.steps_per_epoch == 0 - and self.global_step < self.max_train_steps - ) - ) - if should_checkpoint: + if self.global_step % int(self.args.logging.checkpointing_steps) == 0: self.save_and_eval_checkpoint() if self.suspend_controller.requested(): diff --git a/deepspec/trainer/ckpt_manager.py b/deepspec/trainer/ckpt_manager.py index d04a92f..8b41f86 100644 --- a/deepspec/trainer/ckpt_manager.py +++ b/deepspec/trainer/ckpt_manager.py @@ -12,10 +12,12 @@ from deepspec.utils import ( ensure_dir, + get_rng_state, is_global_main_process, print_on_global_main, print_on_local_main, safe_symlink, + set_rng_state, ) @@ -112,7 +114,9 @@ def load_training_state( assert saved_local_batch_size == int(local_batch_size) torch.set_rng_state(checkpoint["torch_rng"]) - _set_device_rng_state(checkpoint.get("torch_cuda_rng")) + rng_state = checkpoint.get("torch_accelerator_rng", checkpoint.get("torch_cuda_rng")) + if rng_state is not None: + set_rng_state(rng_state) np.random.set_state(checkpoint["numpy_rng"]) random.setstate(checkpoint["python_rng"]) @@ -185,34 +189,6 @@ def save_checkpoint( return checkpoint_dir -def _get_device_rng_state(): - """Return RNG state of the active accelerator (NPU if available, else CUDA).""" - try: - import torch_npu # noqa: F401 - if torch.npu.is_available(): - return torch.npu.get_rng_state() - except Exception: - pass - if torch.cuda.is_available(): - return torch.cuda.get_rng_state() - return None - - -def _set_device_rng_state(state): - """Restore RNG state on the active accelerator (NPU if available, else CUDA).""" - if state is None: - return - try: - import torch_npu # noqa: F401 - if torch.npu.is_available(): - torch.npu.set_rng_state(state) - return - except Exception: - pass - if torch.cuda.is_available(): - torch.cuda.set_rng_state(state) - - def _rank_training_state_path(checkpoint_dir: str, global_rank: int) -> str: return os.path.join( checkpoint_dir, @@ -241,7 +217,7 @@ def _serialize_training_state( "world_size": int(world_size), "local_batch_size": int(local_batch_size), "torch_rng": torch.get_rng_state(), - "torch_cuda_rng": _get_device_rng_state(), + "torch_accelerator_rng": get_rng_state(), "numpy_rng": np.random.get_state(), "python_rng": random.getstate(), } diff --git a/deepspec/utils/__init__.py b/deepspec/utils/__init__.py index c1145d1..6ca7319 100644 --- a/deepspec/utils/__init__.py +++ b/deepspec/utils/__init__.py @@ -4,6 +4,20 @@ import torch from .config import CustomJSONEncoder, jsonable, load_config, parse_opts_to_config +from .device import ( + accelerator_backend, + accelerator_module, + current_device_index, + device_count, + device_type, + empty_cache, + get_rng_state, + is_npu_available, + make_device, + manual_seed_all, + set_device, + set_rng_state, +) from .distributed import ( StatelessResumableDistributedSampler, init_dist, @@ -21,14 +35,7 @@ def seed_all(seed): torch.manual_seed(seed) random.seed(seed) np.random.seed(seed) - try: - import torch_npu # noqa: F401 - torch.npu.manual_seed_all(seed) - return - except Exception: - pass - if torch.cuda.is_available(): - torch.cuda.manual_seed_all(seed) + manual_seed_all(seed) def get_git_sha(detail_info=False): import subprocess @@ -74,21 +81,33 @@ def get_git_diff(rev="HEAD"): "BF16Optimizer", "CustomJSONEncoder", "StatelessResumableDistributedSampler", + "accelerator_backend", + "accelerator_module", "add_metric", + "current_device_index", + "device_count", + "device_type", + "empty_cache", "ensure_dir", "flush", "get_git_diff", "get_git_sha", + "get_rng_state", + "is_npu_available", "init_dist", "is_global_main_process", "is_local_main_process", "jsonable", "load_config", + "make_device", + "manual_seed_all", "main_process_first", "parse_opts_to_config", "print_on_global_main", "print_on_local_main", "reset", "safe_symlink", + "set_device", + "set_rng_state", "seed_all", ] diff --git a/deepspec/utils/device.py b/deepspec/utils/device.py new file mode 100644 index 0000000..87aec33 --- /dev/null +++ b/deepspec/utils/device.py @@ -0,0 +1,76 @@ +import os + +import torch + + +def _npu_module(): + module = getattr(torch, "npu", None) + if module is None: + return None + try: + return module if module.is_available() else None + except Exception: + return None + + +def is_npu_available() -> bool: + requested = os.environ.get("DEEPSPEC_DEVICE", "").strip().lower() + if requested == "cuda": + return False + if requested == "npu": + return _npu_module() is not None + return _npu_module() is not None + + +def device_type() -> str: + requested = os.environ.get("DEEPSPEC_DEVICE", "").strip().lower() + if requested in {"cuda", "npu"}: + return requested + if _npu_module() is not None: + return "npu" + return "cuda" + + +def accelerator_module(): + return getattr(torch, device_type()) + + +def accelerator_backend() -> str: + return "hccl" if device_type() == "npu" else "nccl" + + +def device_count() -> int: + return int(accelerator_module().device_count()) + + +def set_device(local_rank: int) -> None: + accelerator_module().set_device(int(local_rank)) + + +def current_device_index() -> int: + return int(accelerator_module().current_device()) + + +def make_device(local_rank: int | None = None) -> torch.device: + if local_rank is None: + local_rank = current_device_index() + return torch.device(device_type(), int(local_rank)) + + +def manual_seed_all(seed: int) -> None: + if device_count() > 0: + accelerator_module().manual_seed_all(int(seed)) + + +def empty_cache() -> None: + module = accelerator_module() + if hasattr(module, "empty_cache"): + module.empty_cache() + + +def get_rng_state(): + return accelerator_module().get_rng_state() + + +def set_rng_state(state) -> None: + accelerator_module().set_rng_state(state) diff --git a/deepspec/utils/distributed.py b/deepspec/utils/distributed.py index a9193db..3a9c229 100644 --- a/deepspec/utils/distributed.py +++ b/deepspec/utils/distributed.py @@ -8,54 +8,36 @@ import torch.distributed as dist from torch.utils.data import Sampler - -def _is_npu_available(): - try: - import torch_npu # noqa: F401 - return torch.npu.is_available() - except Exception: - return False - - -def _device_backend_info(): - """Return (device_module, device_count_fn, set_device_fn, current_device_fn, backend).""" - if _is_npu_available(): - return ( - torch.npu, - torch.npu.device_count, - torch.npu.set_device, - torch.npu.current_device, - "hccl", - ) - return ( - torch.cuda, - torch.cuda.device_count, - torch.cuda.set_device, - torch.cuda.current_device, - "nccl", - ) +from .device import ( + accelerator_backend, + current_device_index, + device_count, + make_device, + set_device, +) def init_dist(local_rank: int, timeout_minutes: int = 60): - dev_mod, device_count, set_device, _current_device, backend = _device_backend_info() local_world_size = device_count() + assert local_world_size > 0, "no accelerator devices are visible" node_rank = int(os.environ["RANK"]) node_world_size = int(os.environ["WORLD_SIZE"]) rank = node_rank * local_world_size + local_rank world_size = node_world_size * local_world_size init_method = f"tcp://{os.environ['MASTER_ADDR']}:{os.environ['MASTER_PORT']}" set_device(local_rank) - device_type = "npu" if _is_npu_available() else "cuda" - device = torch.device(device_type, local_rank) + device = make_device(local_rank) - dist.init_process_group( - backend=backend, + init_kwargs = dict( + backend=accelerator_backend(), init_method=init_method, rank=rank, world_size=world_size, timeout=timedelta(minutes=timeout_minutes), - device_id=device, ) + if device.type == "cuda": + init_kwargs["device_id"] = device + dist.init_process_group(**init_kwargs) return device, rank, world_size @@ -64,8 +46,7 @@ def is_global_main_process(): def is_local_main_process(): - _dev_mod, _dc, _sd, current_device, _be = _device_backend_info() - return current_device() == 0 + return current_device_index() == 0 def print_on_global_main(*args, **kwargs): diff --git a/deepspec/utils/metrics.py b/deepspec/utils/metrics.py index 5f7eb38..4f8981f 100644 --- a/deepspec/utils/metrics.py +++ b/deepspec/utils/metrics.py @@ -3,6 +3,8 @@ import torch import torch.distributed as dist +from .device import make_device + _REDUCTION_PATTERN = re.compile(r"^(dp_)?(mean|sum|max|min|last)$") _DEFAULT_RATIO_REDUCTION = "dp_sum" @@ -19,8 +21,8 @@ def _detach_scalar(value): def _clone_to_reduce_device(value: torch.Tensor) -> torch.Tensor: tensor = value.detach().clone().to(torch.float32) - if dist.get_backend() == "nccl" and not tensor.is_cuda: - tensor = tensor.to(torch.device("cuda", torch.cuda.current_device())) + if dist.get_backend() in {"nccl", "hccl"} and tensor.device.type == "cpu": + tensor = tensor.to(make_device()) return tensor diff --git a/eval.py b/eval.py index ecd5261..21a2ad8 100644 --- a/eval.py +++ b/eval.py @@ -1,13 +1,11 @@ from __future__ import annotations import argparse import json -import os - import torch from transformers import AutoConfig from deepspec.eval.dspark import Gemma4DSparkEvaluator, Qwen3DSparkEvaluator from deepspec.eval.eagle3 import Gemma4Eagle3Evaluator, Qwen3Eagle3Evaluator -from deepspec.utils import CustomJSONEncoder +from deepspec.utils import CustomJSONEncoder, device_count EVALUATORS = { "Qwen3DSparkModel": Qwen3DSparkEvaluator, @@ -18,7 +16,7 @@ } TASKS = [ - ("gsm8k", 10), + ("gsm8k", 500), ("math500", 500), ("aime25",30), ("humaneval", 164), @@ -28,31 +26,7 @@ ("alpaca", 500), ("arena-hard-v2", 500), ] - -def _is_npu_available(): - try: - import torch_npu # noqa: F401 - return torch.npu.is_available() - except Exception: - return False - - -def _visible_device_count(): - """Return the number of devices actually usable by this process. - - Honors ``ASCEND_RT_VISIBLE_DEVICES`` (NPU) and ``CUDA_VISIBLE_DEVICES`` - (CUDA) so that ``torch.multiprocessing.spawn`` does not over-spawn - workers when the process is restricted to a subset of physical devices. - """ - if _is_npu_available(): - vis = os.environ.get("ASCEND_RT_VISIBLE_DEVICES", "").strip() - if vis: - return len([d for d in vis.split(",") if d.strip() != ""]) - return torch.npu.device_count() - vis = os.environ.get("CUDA_VISIBLE_DEVICES", "").strip() - if vis: - return len([d for d in vis.split(",") if d.strip() != ""]) - return torch.cuda.device_count() +TASK_DEFAULTS = dict(TASKS) def parse_args(): @@ -71,18 +45,28 @@ def parse_args(): parser.add_argument("--step", type=int, default=None,help=("step for tensorboard logging"),) parser.add_argument("--seed", type=int, default=980406) parser.add_argument( - "--tasks", - type=str, + "--task", + action="append", + default=None, + help="Dataset name under eval_datasets without .jsonl. Repeat to run several tasks.", + ) + parser.add_argument( + "--num-samples", + type=int, default=None, - help=("Comma-separated dataset names to evaluate (e.g. gsm8k). Defaults to all tasks."), + help="Override the sample count for every selected task.", ) args = parser.parse_args() - if args.tasks is not None: - selected = {name.strip() for name in args.tasks.split(",")} - args.tasks = [task for task in TASKS if task[0] in selected] - assert args.tasks, f"No matching tasks for --tasks={args.tasks}" - else: - args.tasks = list(TASKS) + task_names = args.task if args.task else [name for name, _ in TASKS] + args.tasks = [ + ( + name, + args.num_samples if args.num_samples is not None else TASK_DEFAULTS.get(name), + ) + for name in task_names + ] + unknown_defaults = [name for name, max_samples in args.tasks if max_samples is None] + assert not unknown_defaults, "--num-samples is required for unknown task(s): " + ", ".join(unknown_defaults) return args @@ -100,5 +84,5 @@ def main(local_rank: int, args): torch.multiprocessing.spawn( main, args=(args,), - nprocs=_visible_device_count(), + nprocs=device_count(), ) diff --git a/requirements.txt b/requirements.txt index 2d69e0f..005ad75 100644 --- a/requirements.txt +++ b/requirements.txt @@ -9,6 +9,7 @@ tensorboard==2.20.0 matplotlib==3.10.9 triton==3.5.1 typing_extensions==4.15.0 +prettytable==3.18.0 sentencepiece==0.2.1 safetensors==0.7.0 diff --git a/scripts/data/prepare_target_cache.py b/scripts/data/prepare_target_cache.py index 744d611..6fa5865 100644 --- a/scripts/data/prepare_target_cache.py +++ b/scripts/data/prepare_target_cache.py @@ -26,6 +26,8 @@ from deepspec.data.jsonl_dataset import JsonLineDataset from deepspec.utils import ( CustomJSONEncoder, + device_count, + empty_cache, get_git_diff, get_git_sha, init_dist, @@ -46,27 +48,6 @@ torch.set_float32_matmul_precision("high") -def _is_npu_available(): - try: - import torch_npu # noqa: F401 - return torch.npu.is_available() - except Exception: - return False - - -def _device_count(): - if _is_npu_available(): - return torch.npu.device_count() - return torch.cuda.device_count() - - -def _empty_cache(): - if _is_npu_available(): - torch.npu.empty_cache() - elif torch.cuda.is_available(): - torch.cuda.empty_cache() - - @dataclass(frozen=True) class TargetForwardResult: target_hidden_states: torch.Tensor @@ -275,7 +256,7 @@ def main(local_rank: int): target_model = AutoModel.from_pretrained( config.model.target_model_name_or_path, dtype=torch.bfloat16, - attn_implementation="eager", + attn_implementation="sdpa", ).to(device=device).eval() target_hidden_size = _get_target_hidden_size(target_model) train_collator = ConversationCollator( @@ -356,7 +337,7 @@ def main(local_rank: int): finally: writer.close() del target_model - _empty_cache() + empty_cache() dataset.close() summary = LocalCacheWriteSummary( global_rank=global_rank, @@ -418,8 +399,6 @@ def main(local_rank: int): if __name__ == "__main__": if os.path.exists(".git"): - try: - print(f"git sha:", get_git_sha()) - except Exception: - pass - torch.multiprocessing.spawn(main, nprocs=_device_count()) + print(f"git status:", "\n\n".join(get_git_sha(detail_info=True))) + print("git diff:", get_git_diff()) + torch.multiprocessing.spawn(main, nprocs=device_count()) diff --git a/scripts/train/train.sh b/scripts/train/train.sh old mode 100644 new mode 100755 index 05a0618..62714dc --- a/scripts/train/train.sh +++ b/scripts/train/train.sh @@ -1,45 +1,48 @@ #!/usr/bin/env bash +set -euo pipefail -# Local launch mirrors the repo's node launcher, not standard -# torchrun semantics. train.py spawns one worker per visible GPU by itself. -# Here RANK/WORLD_SIZE mean node_rank/node_count, so WORLD_SIZE=1 is a -# single-node local run; total GPU workers come from CUDA_VISIBLE_DEVICES. -export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7} +# Prepare the desired Python and accelerator runtime before running this script. +# train.py spawns one worker per visible accelerator. +if [[ "${DEEPSPEC_DEVICE:-}" == "npu" ]]; then + export ASCEND_RT_VISIBLE_DEVICES=${ASCEND_RT_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7} +else + export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7} +fi export MASTER_ADDR=${MASTER_ADDR:-127.0.0.1} export MASTER_PORT=${MASTER_PORT:-29500} export RANK=${RANK:-0} export WORLD_SIZE=${WORLD_SIZE:-1} -# Available public configs: -## dflash -# config/dflash/dflash_gemma4_12b.py -# config/dflash/dflash_qwen3_4b.py -# config/dflash/dflash_qwen3_8b.py -# config/dflash/dflash_qwen3_14b.py -## dspark -# config/dspark/dspark_gemma4_12b.py -# config/dspark/dspark_qwen3_4b.py -# config/dspark/dspark_qwen3_8b.py -# config/dspark/dspark_qwen3_14b.py -## eagle3 -# config/eagle3/eagle3_gemma4_12b.py -# config/eagle3/eagle3_qwen3_4b.py -# config/eagle3/eagle3_qwen3_8b.py -# config/eagle3/eagle3_qwen3_14b.py +config_path=${config_path:-config/dspark/dspark_gemma4_12b.py} +target_model_path=${target_model_path:-google/gemma-4-12B-it} +target_cache_dir=${target_cache_dir:-${HOME}/.cache/deepspec/gemma4_target_cache} +target_layer_ids=${target_layer_ids:-} +global_batch_size=${global_batch_size:-512} +max_train_steps=${max_train_steps:-} +data_max_length=${data_max_length:-4096} +logging_steps=${logging_steps:-10} +checkpointing_steps=${checkpointing_steps:-3000} +exp_name=${exp_name:-} -target_cache_dir=${target_cache_dir:-${HOME}/.cache/deepspec/qwen3_4b_target_cache} - -# --opts overrides any config field by dotted key path: --opts "=". -# Values are parsed as Python scalars (int/float/bool/str). Repeat the flag to set -# multiple fields, e.g.: -# --opts "data.target_cache_path=${target_cache_dir}" \ -# --opts "train.lr=3e-4" \ -# --opts "train.local_batch_size=2" -# -# local_batch_size is the per-GPU micro-batch size. Raise it to better utilize GPUs -# with more memory (e.g. 4 or 8 on 80GB cards), or keep it at 1 if you hit OOM. -# Override it without editing the config via: -# --opts "train.local_batch_size=4" -python train.py \ - --config config/dspark/dspark_qwen3_4b.py \ +cmd=( + python train.py + --config "${config_path}" + --opts "model.target_model_name_or_path=${target_model_path}" --opts "data.target_cache_path=${target_cache_dir}" + --opts "data.max_length=${data_max_length}" + --opts "train.global_batch_size=${global_batch_size}" + --opts "logging.logging_steps=${logging_steps}" + --opts "logging.checkpointing_steps=${checkpointing_steps}" +) + +if [[ -n "${target_layer_ids}" ]]; then + cmd+=(--opts "model.target_layer_ids=${target_layer_ids}") +fi +if [[ -n "${max_train_steps}" ]]; then + cmd+=(--opts "train.max_train_steps=${max_train_steps}") +fi +if [[ -n "${exp_name}" ]]; then + cmd+=(--opts "exp_name=${exp_name}") +fi + +"${cmd[@]}" diff --git a/train.py b/train.py index c8ccfcd..3b8740d 100644 --- a/train.py +++ b/train.py @@ -4,6 +4,7 @@ import torch from deepspec.utils import ( CustomJSONEncoder, + device_count, get_git_diff, load_config, parse_opts_to_config, @@ -17,20 +18,6 @@ torch.set_float32_matmul_precision("high") -def _is_npu_available(): - try: - import torch_npu # noqa: F401 - return torch.npu.is_available() - except Exception: - return False - - -def _device_count(): - if _is_npu_available(): - return torch.npu.device_count() - return torch.cuda.device_count() - - def parse_args(): parser = argparse.ArgumentParser() parser.add_argument("--config", required=True) @@ -54,8 +41,6 @@ def main(local_rank): if __name__ == "__main__": if os.path.exists(".git"): - try: - print(f"git sha:", get_git_sha()) - except Exception: - pass - torch.multiprocessing.spawn(main, nprocs=_device_count()) + print(f"git status:", "\n\n".join(get_git_sha(detail_info=True))) + print("git diff:", get_git_diff()) + torch.multiprocessing.spawn(main, nprocs=device_count()) From 00a7ab841fffd78638ae89a626a1fe3a2b47b6e8 Mon Sep 17 00:00:00 2001 From: hazel duan Date: Mon, 29 Jun 2026 19:34:55 +0800 Subject: [PATCH 3/4] fix: defer Eagle3 trainer imports --- deepspec/trainer/__init__.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/deepspec/trainer/__init__.py b/deepspec/trainer/__init__.py index 5ed28b0..9ccf797 100644 --- a/deepspec/trainer/__init__.py +++ b/deepspec/trainer/__init__.py @@ -1,6 +1,17 @@ from .base_trainer import BaseTrainer from .dspark_trainer import Gemma4DSparkTrainer, Qwen3DSparkTrainer -from .eagle3_trainer import Gemma4Eagle3Trainer, Qwen3Eagle3Trainer + + +def __getattr__(name): + if name in {"Gemma4Eagle3Trainer", "Qwen3Eagle3Trainer"}: + from .eagle3_trainer import Gemma4Eagle3Trainer, Qwen3Eagle3Trainer + + return { + "Gemma4Eagle3Trainer": Gemma4Eagle3Trainer, + "Qwen3Eagle3Trainer": Qwen3Eagle3Trainer, + }[name] + raise AttributeError(name) + __all__ = [ "BaseTrainer", From 72fe3225d59b68973a90f74ec1652cb1ed172745 Mon Sep 17 00:00:00 2001 From: hazel duan Date: Mon, 29 Jun 2026 20:45:01 +0800 Subject: [PATCH 4/4] fix: trim review-only eval changes --- deepspec/eval/dspark/confidence_head.py | 57 +++++++++++++------------ eval.py | 25 +---------- 2 files changed, 30 insertions(+), 52 deletions(-) diff --git a/deepspec/eval/dspark/confidence_head.py b/deepspec/eval/dspark/confidence_head.py index 5380a68..0bddc05 100644 --- a/deepspec/eval/dspark/confidence_head.py +++ b/deepspec/eval/dspark/confidence_head.py @@ -574,34 +574,35 @@ def build_table(self) -> str: int(entry["position"]): entry for entry in row.get("per_position") or [] } - table_row = [ - row["dataset"], - draft_name, - row["sample_count"], - row["proposal_count"], - format_float(summary["ece_mean"]), - format_float(summary["auc_mean"]), - format_float(summary["brier_mean"]), - format_float(summary["pred_mean"]), - format_float(summary["target_mean"]), - ] + [ - ( - format_float(per_position[pos]["ece"]) - if pos in per_position - and float(per_position[pos]["total_weight"]) > 0.0 - else "-" - ) - for pos in range(max_position_count) - ] + [ - ( - format_float(per_position[pos]["auc"]) - if pos in per_position - and float(per_position[pos]["total_weight"]) > 0.0 - else "-" - ) - for pos in range(max_position_count) - ] - table.add_row(table_row) + table.add_row( + [ + row["dataset"], + draft_name, + row["sample_count"], + row["proposal_count"], + format_float(summary["ece_mean"]), + format_float(summary["auc_mean"]), + format_float(summary["brier_mean"]), + format_float(summary["pred_mean"]), + format_float(summary["target_mean"]), + ] + [ + ( + format_float(per_position[pos]["ece"]) + if pos in per_position + and float(per_position[pos]["total_weight"]) > 0.0 + else "-" + ) + for pos in range(max_position_count) + ] + [ + ( + format_float(per_position[pos]["auc"]) + if pos in per_position + and float(per_position[pos]["total_weight"]) > 0.0 + else "-" + ) + for pos in range(max_position_count) + ] + ) return table.get_string() def print_results(self) -> None: diff --git a/eval.py b/eval.py index 21a2ad8..2f630ec 100644 --- a/eval.py +++ b/eval.py @@ -26,8 +26,6 @@ ("alpaca", 500), ("arena-hard-v2", 500), ] -TASK_DEFAULTS = dict(TASKS) - def parse_args(): parser = argparse.ArgumentParser() @@ -44,29 +42,8 @@ def parse_args(): parser.add_argument("--tensorboard-dir", type=str, default=None) parser.add_argument("--step", type=int, default=None,help=("step for tensorboard logging"),) parser.add_argument("--seed", type=int, default=980406) - parser.add_argument( - "--task", - action="append", - default=None, - help="Dataset name under eval_datasets without .jsonl. Repeat to run several tasks.", - ) - parser.add_argument( - "--num-samples", - type=int, - default=None, - help="Override the sample count for every selected task.", - ) args = parser.parse_args() - task_names = args.task if args.task else [name for name, _ in TASKS] - args.tasks = [ - ( - name, - args.num_samples if args.num_samples is not None else TASK_DEFAULTS.get(name), - ) - for name in task_names - ] - unknown_defaults = [name for name, max_samples in args.tasks if max_samples is None] - assert not unknown_defaults, "--num-samples is required for unknown task(s): " + ", ".join(unknown_defaults) + args.tasks = list(TASKS) return args