tacit/nanochat
1
0
Fork 0
mirror of https://github.com/karpathy/nanochat.git synced 2026-07-06 21:09:15 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

fallback to cpu on compute_init function

Browse Source 
fallback to cpu on compute_init function
This commit is contained in:
Sermet Pekin 2025-10-20 11:43:47 +03:00 committed by GitHub
parent 11e46b6439
commit cdb5a455ee
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 10 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

28
nanochat/common.py
View File

@ -91,34 +91,26 @@ def get_dist_info():
def compute_init(): def compute_init():
"""Basic initialization that we keep doing over and over, so make common.""" """Basic initialization that we keep doing over and over, so make common."""
# Check if CUDA is available, otherwise fall back to CPU
# CUDA is currently required if torch.cuda.is_available():
assert torch.cuda.is_available(), "CUDA is needed for a distributed run atm" device = torch.device("cuda")
torch.manual_seed(42)
# Reproducibility torch.cuda.manual_seed(42)
torch.manual_seed(42) else:
torch.cuda.manual_seed(42) device = torch.device("cpu")
# skipping full reproducibility for now, possibly investigate slowdown later torch.manual_seed(42)
# torch.use_deterministic_algorithms(True) logger.warning("CUDA is not available. Falling back to CPU.")
# torch.backends.cudnn.deterministic = True
# torch.backends.cudnn.benchmark = False
# Precision # Precision
torch.set_float32_matmul_precision("high") # uses tf32 instead of fp32 for matmuls torch.set_float32_matmul_precision("high") # uses tf32 instead of fp32 for matmuls
# Distributed setup: Distributed Data Parallel (DDP), optional # Distributed setup: Distributed Data Parallel (DDP), optional
ddp, ddp_rank, ddp_local_rank, ddp_world_size = get_dist_info() ddp, ddp_rank, ddp_local_rank, ddp_world_size = get_dist_info()
if ddp: if ddp and torch.cuda.is_available():
device = torch.device("cuda", ddp_local_rank) device = torch.device("cuda", ddp_local_rank)
torch.cuda.set_device(device) # make "cuda" default to this device torch.cuda.set_device(device) # make "cuda" default to this device
dist.init_process_group(backend="nccl", device_id=device) dist.init_process_group(backend="nccl", device_id=device)
dist.barrier() dist.barrier()
else:
device = torch.device("cuda")
if ddp_rank == 0: if ddp_rank == 0:
logger.info(f"Distributed world size: {ddp_world_size}") logger.info(f"Distributed world size: {ddp_world_size}")
return ddp, ddp_rank, ddp_local_rank, ddp_world_size, device return ddp, ddp_rank, ddp_local_rank, ddp_world_size, device
def compute_cleanup(): def compute_cleanup():