add npu support

requirements-npu.txt

@@ -0,0 +1,7 @@
+torch <= 2.1.0
+torchvision <= 0.16.0
+attrs
+decortator
+scipy
+psutil
+torch_npu

src/training/distributed.py

@@ -1,6 +1,11 @@
 import os
 
 import torch
+try:
+    import torch_npu
+except ImportError:
+    torch_npu = None
+
 import torch.distributed as dist
 
 try:
@@ -107,6 +112,10 @@ def init_distributed_device(args):
         else:
             device = 'cuda:0'
         torch.cuda.set_device(device)
+    if torch_npu != None and torch.npu.is_available():
+        # TODO: add distributed code for npu
+        device = "npu:0"
+        torch_npu.npu.set_device(device)
     else:
         device = 'cpu'
     args.device = device

src/training/main.py

@@ -11,7 +11,10 @@
 import numpy as np
 import torch
 from torch import optim
-from torch.cuda.amp import GradScaler
+try:
+    import torch_npu
+except ImportError:
+    torch_npu = None
 
 try:
     import wandb
@@ -329,6 +332,11 @@ def main(args):
             hvd.broadcast_parameters(model.state_dict(), root_rank=0)
             hvd.broadcast_optimizer_state(optimizer, root_rank=0)
 
+        if args.precision == "amp":
+            if torch_npu != None and torch.npu.is_available():
+                from torch.npu.amp import GradScaler
+            else:
+                from torch.cuda.amp import GradScaler
         scaler = GradScaler() if args.precision == "amp" else None
 
     # optionally resume from a checkpoint

src/training/precision.py

@@ -1,6 +1,9 @@
 import torch
 from contextlib import suppress
-
+try:
+    import torch_npu
+except ImportError:
+    torch_npu = None
 
 def get_autocast(precision):
     if precision == 'amp':

src/training/profiler.py

@@ -1,6 +1,11 @@
 import argparse
 
 import torch
+try:
+    import torch_npu
+except ImportError:
+    torch_npu = None
+
 import open_clip
 import pandas as pd
 from torch.utils.flop_counter import FlopCounterMode
@@ -133,6 +138,8 @@ def profile_model(model_name, batch_size=1, profiler='torch'):
     model.eval()
     if torch.cuda.is_available():
         model = model.cuda()
+    elif torch_npu != None and torch.npu.is_available():
+        model = model.npu()
 
     if isinstance(model.visual.image_size, (tuple, list)):
         image_input_size = (3,) + tuple(model.visual.image_size[-2:])