merged single gpu changes and some fixes

Author: ShashankSubramanian

.gitignore

@@ -158,3 +158,6 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+# output logs
+*.out

README.md

@@ -3,25 +3,26 @@ base: &base
   # Model config
   embed_dim: 384
   depth: 12
-  dropout: 0
+  dropout: 0.0
   patch_size: 8
   num_heads: 8
 
   # Training config
   img_size: [360, 720]
   dt: 1 
-  global_batch_size: 256 # number of samples per training batch
-  num_epochs: 60
+  global_batch_size: 16 # number of samples per training batch
+  num_epochs: 25
   amp_mode: none
-  enable_apex: False
-  enable_jit: False
+  enable_fused: false
+  enable_jit: false
   expdir: '/logs'
   lr_schedule: 'cosine'
   lr: 5E-4
+  warmup: 0
 
   # Data
   data_loader_config: 'pytorch'
-  num_data_workers: 2 # number of dataloader worker threads per proc
+  num_data_workers: 0 # number of dataloader worker threads per proc
   n_in_channels: 20
   n_out_channels: 20
   train_data_path:   '/data/train'
@@ -30,28 +31,52 @@ base: &base
   time_means_path:   '/data/stats/time_means.npy'
   global_means_path: '/data/stats/global_means.npy'
   global_stds_path:  '/data/stats/global_stds.npy'
+  limit_nsamples: None
+  limit_nsamples_val: None
 
   # Comms
   wireup_info: env
   wireup_store: tcp
 
-short_noopt:
-  <<: *base
-  num_epochs: 10
-  num_data_workers: 0
-  global_batch_size: 64
-
 short: &short
   <<: *base
   num_epochs: 10
   num_data_workers: 8
   global_batch_size: 64
   embed_dim: 1024
 
+short_limitsamples: &short_ls
+  <<: *base
+  limit_nsamples: 512
+  limit_nsamples_val: 128
+  num_epochs: 4
+
 # Short config with full optimizations
 short_opt:
-  <<: *short
-  global_batch_size: 16
+  <<: *short_ls
+  global_batch_size: 64
+  data_loader_config: 'dali'
+  num_data_workers: 8
+  amp_mode: bf16
+  enable_jit: true
+  enable_fused: true
+
+
+# Model parallel configs
+short_mp: &short_mp
+  <<: *short_ls
+  global_batch_size: 64
   data_loader_config: 'dali'
   num_data_workers: 8
   embed_dim: 1024
+  amp_mode: bf16
+  enable_jit: true
+  enable_fused: true
+
+mp:
+  <<: *short_mp
+  global_batch_size: 64
+  limit_nsamples: None
+  limit_nsamples_val: None
+  num_epochs: 20
+  data_loader_config: 'dali'

config/ViT.yaml

@@ -9,14 +9,12 @@ mkdir -p ${LOGDIR}
 
 ngpu=4
 config_file=./config/ViT.yaml
-config="short_opt"
-run_num="testgraphs"
-amp_mode="fp16"
+config="short_mp"
+run_num="0"
 col_parallel_size=1
-row_parallel_size=2
-local_batch_size=16
+row_parallel_size=1
 #cmd="python train.py --amp_mode=$amp_mode --yaml_config=$config_file --config=$config --run_num=$run_num --local_batch_size=$local_batch_size"
-cmd="python train_mp_graphs.py --local_batch_size=$local_batch_size --row_parallel_size=$row_parallel_size --col_parallel_size=$col_parallel_size --amp_mode=$amp_mode --yaml_config=$config_file --config=$config --run_num=$run_num"
+cmd="python train_mp.py --row_parallel_size=$row_parallel_size --col_parallel_size=$col_parallel_size --yaml_config=$config_file --config=$config --run_num=$run_num"
 
 
 srun -n $ngpu --cpus-per-task=32 --gpus-per-node $ngpu shifter --image=${image} --module=gpu,nccl-2.18 -V ${DATADIR}:/data -V ${LOGDIR}:/logs  bash -c "source export_DDP_vars.sh && $cmd"

interactive_train.sh

@@ -38,7 +38,7 @@
    },
    "outputs": [],
    "source": [
-    "log_dir = os.path.expandvars('${SCRATCH}/sc23-dl-tutorial/logs/short_opt/4GPU/bs16_graphs')"
+    "log_dir = os.path.expandvars('${SCRATCH}/sc23-dl-tutorial/logs/short_mp/4MP/bs64')"
    ]
   },
   {
@@ -65,7 +65,7 @@
     {
      "data": {
       "text/html": [
-       "<a href=\"https://jupyter.nersc.gov/user/shas1693/perlmutter-login-node-base/proxy/33151/\">https://jupyter.nersc.gov/user/shas1693/perlmutter-login-node-base/proxy/33151/</a>"
+       "<a href=\"https://jupyter.nersc.gov/user/shas1693/perlmutter-login-node-base/proxy/35199/\">https://jupyter.nersc.gov/user/shas1693/perlmutter-login-node-base/proxy/35199/</a>"
       ],
       "text/plain": [
        "<IPython.core.display.HTML object>"

sample_nsys_profiles/4workers.nsys-rep

@@ -1,34 +1,27 @@
 #!/bin/bash 
-#SBATCH -C gpu 
-#SBATCH --nodes=1
-#SBATCH -q regular
-#SBATCH -A nstaff
+#SBATCH -C 'gpu&hbm80g'
 #SBATCH --ntasks-per-node 4
 #SBATCH --cpus-per-task 32
 #SBATCH --gpus-per-node 4
-#SBATCH --time=03:00:00
+#SBATCH --time=06:00:00
 #SBATCH --image=nersc/pytorch:ngc-23.04-v0
 #SBATCH --module=gpu,nccl-2.18
-#SBATCH -J vit-era5
+#SBATCH -J vit-era5-dp-ng
 #SBATCH -o %x-%j.out
 
 DATADIR=/pscratch/sd/s/shas1693/data/sc23_tutorial_data/downsampled
 LOGDIR=${SCRATCH}/sc23-dl-tutorial/logs
 mkdir -p ${LOGDIR}
 
 config_file=./config/ViT.yaml
-config="short_opt"
-run_num="bs16_graphs"
-suffix="_graphs"
-#suffix=""
-amp_mode="fp16"
+config="mp"
+run_num="bs64"
+suffix=""
 col_parallel_size=1
-row_parallel_size=4
-args="--col_parallel_size=$col_parallel_size --row_parallel_size=$row_parallel_size --amp_mode=$amp_mode --yaml_config=$config_file --config=$config --run_num=$run_num"
-
+row_parallel_size=1
+args="--col_parallel_size=$col_parallel_size --row_parallel_size=$row_parallel_size --yaml_config=$config_file --config=$config --run_num=$run_num"
 
 export FI_MR_CACHE_MONITOR=userfaultfd
-export NCCL_NET_GDR_LEVEL=PHB
 export HDF5_USE_FILE_LOCKING=FALSE
 
 # Profiling

sample_nsys_profiles/baseline.nsys-rep

@@ -3,7 +3,6 @@
 from utils.YParams import YParams
 from torchinfo import summary
 
-params = YParams('./config/ViT.yaml', 'base')
-params.device = 'gpu'
+params = YParams('./config/ViT.yaml', 'short_mp')
 model = ViT(params)
-summary(model, input_size=(1,20,360,720))
+summary(model, input_size=(16,20,360,720))

sample_nsys_profiles/dali.nsys-rep

@@ -19,21 +19,9 @@
 from utils import get_data_loader_distributed
 from utils.loss import l2_loss, l2_loss_opt
 from utils.metrics import weighted_rmse
+from utils.plots import generate_images
 from networks import vit
 
-import apex.optimizers as aoptim
-
-def compute_grad_norm(p_list, device):
-    norm_type = 2.0
-    grads = [p.grad for p in p_list if p.grad is not None]
-    total_norm = torch.norm(torch.stack([torch.norm(g.detach(), norm_type).to(device) for g in grads]), norm_type)
-    return total_norm
-
-def compute_parameter_norm(p_list, device):
-    norm_type = 2.0
-    total_norm = torch.norm(torch.stack([torch.norm(p.detach(), norm_type).to(device) for p in p_list]), norm_type)
-    return total_norm
-
 def train(params, args, local_rank, world_rank, world_size):
     # set device and benchmark mode
     torch.backends.cudnn.benchmark = True
@@ -48,6 +36,9 @@ def train(params, args, local_rank, world_rank, world_size):
 
     # create model
     model = vit.ViT(params).to(device)
+
+    if params.enable_jit:
+        model = torch.compile(model)
 
     if params.amp_dtype == torch.float16: 
         scaler = GradScaler()
@@ -61,9 +52,8 @@ def train(params, args, local_rank, world_rank, world_size):
             model = DistributedDataParallel(model, device_ids=[local_rank],
                                             bucket_cap_mb=args.bucket_cap_mb)
 
-    if params.enable_apex:
-        optimizer = aoptim.FusedAdam(model.parameters(), lr = params.lr,
-                                    adam_w_mode=False, set_grad_none=True)
+    if params.enable_fused:
+        optimizer = optim.Adam(model.parameters(), lr = params.lr, fused=True)
     else:
         optimizer = optim.Adam(model.parameters(), lr = params.lr)
 
@@ -74,14 +64,14 @@ def train(params, args, local_rank, world_rank, world_size):
     startEpoch = 0
 
     if params.lr_schedule == 'cosine':
-        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=params.num_epochs, last_epoch=startEpoch-1)
+        if params.warmup > 0:
+            lr_scale = lambda x: min(params.lr*((x+1)/params.warmup), 0.5*params.lr*(1 + np.cos(np.pi*x/params.num_epochs)))
+            scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_scale)
+        else:
+            scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=params.num_epochs, last_epoch=startEpoch-1)
     else:
         scheduler = None
 
-    if params.enable_jit:
-        model_handle = model.module if (params.distributed and not args.noddp) else model
-        model_handle = torch.jit.script(model_handle)  
-
     # select loss function
     if params.enable_jit:
         loss_func = l2_loss_opt
@@ -124,51 +114,48 @@ def train(params, args, local_rank, world_rank, world_size):
         model.train()
         step_count = 0
         for i, data in enumerate(train_data_loader, 0):
-            if (args.enable_manual_profiling and world_rank==0):
+            if world_rank == 0:
                 if (epoch == 3 and i == 0):
                     torch.cuda.profiler.start()
-                if (epoch == 3 and i == 59):
+                if (epoch == 3 and i == len(train_data_loader) - 1):
                     torch.cuda.profiler.stop()
 
-            if args.enable_manual_profiling: torch.cuda.nvtx.range_push(f"step {i}")
+            torch.cuda.nvtx.range_push(f"step {i}")
             iters += 1
             dat_start = time.time()
-            if args.enable_manual_profiling: torch.cuda.nvtx.range_push(f"data copy in {i}")
+            torch.cuda.nvtx.range_push(f"data copy in {i}")
 
             inp, tar = map(lambda x: x.to(device), data)
-            if args.enable_manual_profiling: torch.cuda.nvtx.range_pop() # copy in
+            torch.cuda.nvtx.range_pop() # copy in
 
             tr_start = time.time()
             b_size = inp.size(0)
 
             optimizer.zero_grad()
 
-            if args.enable_manual_profiling: torch.cuda.nvtx.range_push(f"forward")
+            torch.cuda.nvtx.range_push(f"forward")
             with autocast(enabled=params.amp_enabled, dtype=params.amp_dtype):
                 gen = model(inp)
                 loss = loss_func(gen, tar)
-            if args.enable_manual_profiling: torch.cuda.nvtx.range_pop() #forward
+            torch.cuda.nvtx.range_pop() #forward
 
             if params.amp_dtype == torch.float16: 
                 scaler.scale(loss).backward()
-                if args.enable_manual_profiling: torch.cuda.nvtx.range_push(f"optimizer")
+                torch.cuda.nvtx.range_push(f"optimizer")
                 scaler.step(optimizer)
-                if args.enable_manual_profiling: torch.cuda.nvtx.range_pop() # optimizer
+                torch.cuda.nvtx.range_pop() # optimizer
                 scaler.update()
             else:
                 loss.backward()
-                if args.enable_manual_profiling: torch.cuda.nvtx.range_push(f"optimizer")
+                torch.cuda.nvtx.range_push(f"optimizer")
                 optimizer.step()
-                if args.enable_manual_profiling: torch.cuda.nvtx.range_pop() # optimizer
+                torch.cuda.nvtx.range_pop() # optimizer
 
             if params.distributed:
                 torch.distributed.all_reduce(loss)
             tr_loss.append(loss.item()/world_size)
 
-            if args.enable_manual_profiling: torch.cuda.nvtx.range_pop() # step
-
-#            g_norm = compute_grad_norm(model.parameters(), device)
-#            p_norm = compute_parameter_norm(model.parameters(), device)
+            torch.cuda.nvtx.range_pop() # step
 
             tr_end = time.time()
             tr_time += tr_end - tr_start
@@ -187,6 +174,8 @@ def train(params, args, local_rank, world_rank, world_size):
             args.tboard_writer.add_scalar('Loss/train', np.mean(tr_loss), iters)
             args.tboard_writer.add_scalar('Learning Rate', optimizer.param_groups[0]['lr'], iters)
             args.tboard_writer.add_scalar('Avg iters per sec', step_count/(end-start), iters)
+            fig = generate_images([inp, tar, gen])
+            args.tboard_writer.add_figure('Visualization, t2m', fig, iters, close=True)
 
         val_start = time.time()
         val_loss = []
@@ -227,12 +216,12 @@ def train(params, args, local_rank, world_rank, world_size):
     parser.add_argument("--yaml_config", default='./config/ViT.yaml', type=str, help='path to yaml file containing training configs')
     parser.add_argument("--config", default='base', type=str, help='name of desired config in yaml file')
     parser.add_argument("--amp_mode", default='none', type=str, choices=['none', 'fp16', 'bf16'], help='select automatic mixed precision mode')  
-    parser.add_argument("--enable_apex", action='store_true', help='enable apex fused Adam optimizer')
+    parser.add_argument("--enable_fused", action='store_true', help='enable fused Adam optimizer')
     parser.add_argument("--enable_jit", action='store_true', help='enable JIT compilation')
-    parser.add_argument("--enable_manual_profiling", action='store_true', help='enable manual nvtx ranges and profiler start/stop calls')
     parser.add_argument("--local_batch_size", default=None, type=int, help='local batchsize (manually override global_batch_size config setting)')
     parser.add_argument("--num_epochs", default=None, type=int, help='number of epochs to run')
     parser.add_argument("--num_data_workers", default=None, type=int, help='number of data workers for data loader')
+    parser.add_argument("--data_loader_config", default=None, type=str, choices=['pytorch', 'dali'], help="dataloader configuration. choices: 'pytorch', 'dali'")
     parser.add_argument("--bucket_cap_mb", default=25, type=int, help='max message bucket size in mb')
     parser.add_argument("--disable_broadcast_buffers", action='store_true', help='disable syncing broadcasting buffers')
     parser.add_argument("--noddp", action='store_true', help='disable DDP communication')
@@ -253,9 +242,12 @@ def train(params, args, local_rank, world_rank, world_size):
         amp_dtype = torch.bfloat16    
     params.update({"amp_enabled": amp_dtype is not torch.float32,
                     "amp_dtype" : amp_dtype, 
-                    "enable_apex" : args.enable_apex,
+                    "enable_fused" : args.enable_fused,
                     "enable_jit" : args.enable_jit
                     })
+
+    if args.data_loader_config:
+        params.update({"data_loader_config" : args.data_loader_config})
 
     if args.num_epochs:
         params.update({"num_epochs" : args.num_epochs})