diff --git a/deepspec/data/cuda_prefetcher.py b/deepspec/data/cuda_prefetcher.py index 69eac90..7f682a4 100644 --- a/deepspec/data/cuda_prefetcher.py +++ b/deepspec/data/cuda_prefetcher.py @@ -2,6 +2,8 @@ import torch +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()} @@ -22,7 +24,8 @@ class CUDAPrefetcher: def __init__(self, dataloader, device): self.dataloader = dataloader self.device = device - self.stream = torch.cuda.Stream(device=device) + self.accelerator = accelerator_module() + self.stream = self.accelerator.Stream(device=device) def __iter__(self): self._iter = iter(self.dataloader) @@ -40,7 +43,7 @@ def _fetch_and_transfer(self): except StopIteration: self._done = True return - with torch.cuda.stream(self.stream): + with self.accelerator.stream(self.stream): self._gpu_batch = move_batch_to_device(cpu_batch, self.device) def __next__(self): @@ -53,7 +56,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.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 36b1c32..0bddc05 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 = ( @@ -363,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, diff --git a/deepspec/modeling/dspark/common.py b/deepspec/modeling/dspark/common.py index 76d908a..1a6f6fb 100644 --- a/deepspec/modeling/dspark/common.py +++ b/deepspec/modeling/dspark/common.py @@ -82,7 +82,34 @@ def create_dspark_attention_mask( seq_len: int, block_size: int, device: torch.device, + 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 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 ae71fac..03ccf34 100644 --- a/deepspec/modeling/dspark/qwen3/config.py +++ b/deepspec/modeling/dspark/qwen3/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 build_draft_config( diff --git a/deepspec/modeling/dspark/qwen3/modeling.py b/deepspec/modeling/dspark/qwen3/modeling.py index 2f9c220..3c3dd01 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, + attn_implementation=self.config._attn_implementation, ) output_hidden = self._forward_backbone( position_ids=full_position_ids, 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", diff --git a/deepspec/trainer/base_trainer.py b/deepspec/trainer/base_trainer.py index 1a03dc8..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, @@ -65,9 +67,9 @@ def _build_fsdp_kwargs( sharding_strategy=sharding_strategy, ) if sharding_strategy in _HYBRID_STRATEGIES: - devices_per_node = torch.cuda.device_count() + devices_per_node = device_count() fsdp_kwargs["device_mesh"] = init_device_mesh( - "cuda", + device_type(), (world_size // devices_per_node, devices_per_node), mesh_dim_names=("replicate", "shard"), ) @@ -283,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): diff --git a/deepspec/trainer/ckpt_manager.py b/deepspec/trainer/ckpt_manager.py index 6b99f65..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"]) - torch.cuda.set_rng_state(checkpoint["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"]) @@ -213,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": torch.cuda.get_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 cdf9aff..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,7 +35,7 @@ def seed_all(seed): torch.manual_seed(seed) random.seed(seed) np.random.seed(seed) - torch.cuda.manual_seed_all(seed) + manual_seed_all(seed) def get_git_sha(detail_info=False): import subprocess @@ -67,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 868bb2d..3a9c229 100644 --- a/deepspec/utils/distributed.py +++ b/deepspec/utils/distributed.py @@ -8,25 +8,36 @@ import torch.distributed as dist from torch.utils.data import Sampler +from .device import ( + accelerator_backend, + current_device_index, + device_count, + make_device, + set_device, +) + def init_dist(local_rank: int, timeout_minutes: int = 60): - local_world_size = torch.cuda.device_count() + 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']}" - torch.cuda.set_device(local_rank) - device = torch.device("cuda", local_rank) + set_device(local_rank) + device = make_device(local_rank) - dist.init_process_group( - backend="nccl", + 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 @@ -35,7 +46,7 @@ def is_global_main_process(): def is_local_main_process(): - return torch.cuda.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 a35e7ea..2f630ec 100644 --- a/eval.py +++ b/eval.py @@ -5,7 +5,7 @@ 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, @@ -61,5 +61,5 @@ def main(local_rank: int, args): torch.multiprocessing.spawn( main, args=(args,), - nprocs=torch.cuda.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 d072582..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, @@ -335,7 +337,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, @@ -399,4 +401,4 @@ def main(local_rank: int): 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()) + 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 40d9e07..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, @@ -42,4 +43,4 @@ def main(local_rank): 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()) + torch.multiprocessing.spawn(main, nprocs=device_count())