Add torch.accelerator API to mingGPT example (#34)

* Add torch.accelerator API to mingGPT example

Author: gera-aldama

distributed/minGPT-ddp/mingpt/main.py

@@ -1,4 +1,5 @@
 import os
+import sys
 import torch
 from torch.utils.data import random_split
 from torch.distributed import init_process_group, destroy_process_group
@@ -8,10 +9,18 @@
 from omegaconf import DictConfig
 import hydra
 
+def verify_min_gpu_count(min_gpus: int = 2) -> bool:
+ has_gpu = torch.accelerator.is_available()
+ gpu_count = torch.accelerator.device_count()
+ return has_gpu and gpu_count >= min_gpus
 
 def ddp_setup():
- init_process_group(backend="nccl")
- torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
+ acc = torch.accelerator.current_accelerator()
+ rank = int(os.environ["LOCAL_RANK"])
+ device: torch.device = torch.device(f"{acc}:{rank}")
+ backend = torch.distributed.get_default_backend_for_device(device)
+ init_process_group(backend=backend)
+ torch.accelerator.set_device_index(rank)
 
 def get_train_objs(gpt_cfg: GPTConfig, opt_cfg: OptimizerConfig, data_cfg: DataConfig):
  dataset = CharDataset(data_cfg)
@@ -40,6 +49,9 @@ def main(cfg: DictConfig):
 
  destroy_process_group()
 
-
 if __name__ == "__main__":
+ _min_gpu_count = 2
+ if not verify_min_gpu_count(min_gpus=_min_gpu_count):
+ print(f"Unable to locate sufficient {_min_gpu_count} gpus to run this example. Exiting.")
+ sys.exit()
  main()

distributed/minGPT-ddp/mingpt/trainer.py

@@ -48,6 +48,10 @@ def __init__(self, trainer_config: TrainerConfig, model, optimizer, train_datase
  # set torchrun variables
  self.local_rank = int(os.environ["LOCAL_RANK"])
  self.global_rank = int(os.environ["RANK"]) 
+ # set device
+ self.acc = torch.accelerator.current_accelerator()
+ self.device: torch.device = torch.device(f"{self.acc}:{self.local_rank}")
+ self.device_type = self.device.type
  # data stuff
  self.train_dataset = train_dataset
  self.train_loader = self._prepare_dataloader(train_dataset)
@@ -58,7 +62,7 @@ def __init__(self, trainer_config: TrainerConfig, model, optimizer, train_datase
  self.optimizer = optimizer 
  self.save_every = self.config.save_every
  if self.config.use_amp:
- self.scaler = torch.cuda.amp.GradScaler()
+ self.scaler = torch.amp.GradScaler(self.device_type)
  # load snapshot if available. only necessary on the first node.
  if self.config.snapshot_path is None:
  self.config.snapshot_path = "snapshot.pt"
@@ -93,7 +97,7 @@ def _load_snapshot(self):
 
 
  def _run_batch(self, source, targets, train: bool = True) -> float:
- with torch.set_grad_enabled(train), torch.amp.autocast(device_type="cuda", dtype=torch.float16, enabled=(self.config.use_amp)):
+ with torch.set_grad_enabled(train), torch.amp.autocast(device_type=self.device_type, dtype=torch.float16, enabled=(self.config.use_amp)):
  _, loss = self.model(source, targets)
 
  if train:
@@ -119,7 +123,7 @@ def _run_epoch(self, epoch: int, dataloader: DataLoader, train: bool = True):
  targets = targets.to(self.local_rank)
  batch_loss = self._run_batch(source, targets, train)
  if iter % 100 == 0:
- print(f"[GPU{self.global_rank}] Epoch {epoch} | Iter {iter} | {step_type} Loss {batch_loss:.5f}")
+ print(f"[RANK{self.global_rank}] Epoch {epoch} | Iter {iter} | {step_type} Loss {batch_loss:.5f}")
 
  def _save_snapshot(self, epoch):
  # capture snapshot

distributed/minGPT-ddp/requirements.txt

@@ -1,5 +1,4 @@
-torch>=1.11.0
-fsspec
+torch>=2.7
 boto3
 hydra-core
 requests

distributed/minGPT-ddp/run_example.sh

@@ -0,0 +1,10 @@
+#!/bin/bash
+# Usage: bash run_example.sh {file_to_run.py} {num_gpus}
+# where file_to_run = example to run. Default = 'main.py'
+# num_gpus = num local gpus to use. Default = 16
+
+# samples to run include:
+# main.py
+
+echo "Launching ${1:-main.py} with ${2:-16} gpus"
+torchrun --standalone --nproc_per_node=${2:-16} ${1:-main.py}

run_distributed_examples.sh

@@ -54,9 +54,14 @@ function distributed_ddp() {
  uv run main.py || error "ddp example failed"
 }
 
+function distributed_minGPT-ddp() {
+ uv run bash run_example.sh mingpt/main.py || error "minGPT example failed"
+}
+
 function run_all() {
  run distributed/tensor_parallelism
  run distributed/ddp
+ run distributed/minGPT-ddp
 }
 
 # by default, run all examples