17160118

context_unet.py

@@ -32,15 +32,11 @@ class GroupNorm32(nn.GroupNorm):
         self.swish = swish
 
     def forward(self, x):
-        #print(f"GroupNorm32, x.dtype = {x.dtype}, x.float().dtype = {x.float().dtype}, swish = {self.swish}")
-        #y = super().forward(x.float()).to(x.dtype)
         y = super().forward(x)
-        #print(f"swish == {self.swish}, {y.dtype}")
         if self.swish == 1.0:
             y = F.silu(y)
         elif self.swish:
             y = y * F.sigmoid(y * float(self.swish))
-        #print(f"swish == {self.swish}, {y.dtype}")
         return y
 
 def normalization(channels, swish=0.0):
@@ -191,8 +187,7 @@ class ResBlock(TimestepBlock):
             h = in_conv(h)
         else:
             h = self.in_layers(x)
-        # print("forward, h.dtype =", h.dtype)
-        emb_out = self.emb_layers(emb).type(h.dtype)
+        emb_out = self.emb_layers(emb)#.type(h.dtype)
 
         while len(emb_out.shape) < len(h.shape):
             emb_out = emb_out[..., None]
@@ -230,7 +225,7 @@ class QKVAttention(nn.Module):
         scale = 1 / math.sqrt(math.sqrt(ch))
         weight = torch.einsum("bct,bcs->bts", q*scale, k*scale)
         # print("forward, weight.dtype =", weight.dtype)
-        weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype)
+        weight = torch.softmax(weight.float(), dim=-1)#.type(weight.dtype)
 
         a = torch.einsum("bts,bcs->bct", weight, v)
         return a.reshape(bs, -1, length)
@@ -290,7 +285,7 @@ def timestep_embedding(timesteps, dim, max_period=10000):
     #print(f"timestep_embedding is running")
     half = dim // 2
     freqs = torch.exp(
-        -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+        -math.log(max_period) * torch.arange(start=0, end=half) / half #, dtype=torch.float32) / half
     ).to(device=timesteps.device)
     #print (timesteps[:, None].float().shape,freqs[None].shape)
     args = timesteps[:, None].float() * freqs[None]
@@ -322,7 +317,7 @@ class ContextUnet(nn.Module):
         encoder_channels = None,
         dim = 2,
         stride = (2,2),
-        dtype = torch.float32,
+        #dtype = torch.float32,
         ):
         super().__init__()
 
@@ -356,7 +351,7 @@ class ContextUnet(nn.Module):
         # self.n_param = n_param
         self.model_channels = model_channels
         # self.use_fp16 = use_fp16
-        self.dtype = dtype#torch.float16 if self.use_fp16 else torch.float32
+        #self.dtype = dtype#torch.float16 if self.use_fp16 else torch.float32
 
         self.token_embedding = nn.Linear(n_param, model_channels * 4)
 
@@ -526,15 +521,15 @@ class ContextUnet(nn.Module):
     def forward(self, x, timesteps, y=None):
         hs = []
         # print("device of timesteps, self.model_channels:", timesteps.device, self.model_channels)
-        emb = self.time_embed(timestep_embedding(timesteps, self.model_channels).to(self.dtype))
+        emb = self.time_embed(timestep_embedding(timesteps, self.model_channels))#.to(self.dtype))
         #print(f"forward after emb")
         if y != None:
             #text_outputs = self.token_embedding(y.float())
-            text_outputs = self.token_embedding(y.to(self.dtype))
+            text_outputs = self.token_embedding(y)#.to(self.dtype))
             emb = emb + text_outputs.to(emb)
 
         #print("forward, h = x.type(self.dtype), self.dtype =", self.dtype)
-        h = x.type(self.dtype)
+        h = x.clone()#.type(self.dtype)
         #print("0,h.shape =", h.shape)
         for module in self.input_blocks:
             h = module(h, emb)
@@ -552,7 +547,7 @@ class ContextUnet(nn.Module):
             # print("module decoder, h.shape =", h.shape)
 
         #print("h = h.type(x.dtype), x.dtype =", x.dtype, h.dtype)
-        h = h.type(x.dtype)
+        #h = h.type(x.dtype)
         h = self.out(h)
         #print("self.out(h)", "h.dtype =", h.dtype)
 

diffusion.py

@@ -77,6 +77,8 @@ import sys
 from datetime import timedelta
 from time import time
 
+from torch.cuda.amp import autocast, GradScaler
+
 # %%
 def ddp_setup(rank: int, world_size: int, master_addr, master_port):
     """
@@ -117,9 +119,9 @@ def ddp_setup(rank: int, world_size: int, master_addr, master_port):
 
 # %%
 class DDPMScheduler(nn.Module):
-    def __init__(self, betas: tuple, num_timesteps: int, img_shape: list, device='cpu', dtype=torch.float16, config=None):
+    def __init__(self, betas: tuple, num_timesteps: int, img_shape: list, device='cpu', config=None):#, dtype=torch.float16,
         super().__init__()
-        self.dtype = dtype#torch.float16 if self.use_fp16 else torch.float32
+        #self.dtype = dtype#torch.float16 if self.use_fp16 else torch.float32
         
         beta_1, beta_T = betas
         assert 0 < beta_1 <= beta_T <= 1, "ensure 0 < beta_1 <= beta_T <= 1"
@@ -127,7 +129,7 @@ class DDPMScheduler(nn.Module):
         self.num_timesteps = num_timesteps
         self.img_shape = img_shape
         self.beta_t = torch.linspace(beta_1, beta_T, self.num_timesteps) #* (beta_T-beta_1) + beta_1
-        self.beta_t = self.beta_t.to(self.dtype)
+        #self.beta_t = self.beta_t.to(self.dtype)
         self.beta_t = self.beta_t.to(self.device)
 
         # self.drop_prob = drop_prob
@@ -160,7 +162,7 @@ class DDPMScheduler(nn.Module):
     def sample(self, nn_model, params, device, guide_w = 0):
         n_sample = len(params) #params.shape[0]
         # print("params.shape[0], len(params)", params.shape[0], len(params))
-        x_i = torch.randn(n_sample, *self.img_shape).to(self.dtype)
+        x_i = torch.randn(n_sample, *self.img_shape)#.to(self.dtype)
         x_i = x_i.to(device)
         #print(f"#1 x_i.device = {x_i.device}")
         # print("x_i.shape =", x_i.shape)
@@ -171,7 +173,7 @@ class DDPMScheduler(nn.Module):
             # uncond_tokens = torch.tensor(np.float32(np.array([0,0]))).to(device)
             # uncond_tokens = uncond_tokens.repeat(int(n_sample),1)
             #c_i = torch.cat((c_i, uncond_tokens), 0)
-            c_i = c_i.to(self.dtype)
+            #c_i = c_i.to(self.dtype)
 
         x_i_entire = [] # keep track of generated steps in case want to plot something
         # print("self.num_timesteps =", self.num_timesteps)
@@ -183,14 +185,14 @@ class DDPMScheduler(nn.Module):
             # print(f'sampling timestep {i:4d}',end='\r')
             t_is = torch.tensor([i]).to(device)
             t_is = t_is.repeat(n_sample)
-            t_is = t_is.to(self.dtype)
+            #t_is = t_is.to(self.dtype)
 
             z = torch.randn(n_sample, *self.img_shape).to(device) if i > 0 else torch.tensor(0.)
-            z = z.to(self.dtype)
+            #z = z.to(self.dtype)
 
             if guide_w == -1:
                 # eps = nn_model(x_i, t_is, return_dict=False)[0]
-                eps = nn_model(x_i, t_is).to(self.dtype)
+                eps = nn_model(x_i, t_is)#.to(self.dtype)
                 # x_i = 1/torch.sqrt(self.alpha_t[i])*(x_i-eps*self.beta_t[i]/torch.sqrt(1-self.bar_alpha_t[i])) + torch.sqrt(self.beta_t[i])*z
             else:
                 # double batch
@@ -201,7 +203,7 @@ class DDPMScheduler(nn.Module):
                 # split predictions and compute weighting
                 # print("nn_model input shape", x_i.shape, t_is.shape, c_i.shape)
                 #print(f"sample, i = {i}, x_i.dtype = {x_i.dtype}, c_i.dtype = {c_i.dtype}")
-                eps = nn_model(x_i, t_is, c_i).to(self.dtype)
+                eps = nn_model(x_i, t_is, c_i)#.to(self.dtype)
                 #eps1 = eps[:n_sample]
                 #eps2 = eps[n_sample:]
                 #eps = eps1 + guide_w*(eps1 - eps2)
@@ -317,8 +319,8 @@ class TrainConfig:
     # data_dir = './data' # data directory
 
     #use_fp16 = True 
-    dtype = torch.float32 #if use_fp16 else torch.float32
-    mixed_precision = "no" #"fp16"
+    #dtype = torch.float32 #if use_fp16 else torch.float32
+    #mixed_precision = "no" #"fp16"
     gradient_accumulation_steps = 1
 
     pbar_update_step = 20 
@@ -389,11 +391,11 @@ class DDPM21CM:
         # self.dataloader = DataLoader(dataset, batch_size=config.batch_size, shuffle=True)
         # del dataset
         # print("self.ddpm = DDPMScheduler")
-        self.ddpm = DDPMScheduler(betas=(1e-4, 0.02), num_timesteps=config.num_timesteps, img_shape=config.img_shape, device=config.device, dtype=config.dtype, config=config,)
+        self.ddpm = DDPMScheduler(betas=(1e-4, 0.02), num_timesteps=config.num_timesteps, img_shape=config.img_shape, device=config.device, config=config,)#, dtype=config.dtype
 
         # print("self.nn_model = ContextUnet")
         # initialize the unet
-        self.nn_model = ContextUnet(n_param=config.n_param, image_size=config.HII_DIM, dim=config.dim, stride=config.stride, dtype=config.dtype)
+        self.nn_model = ContextUnet(n_param=config.n_param, image_size=config.HII_DIM, dim=config.dim, stride=config.stride)#, dtype=config.dtype)
 
         # print("self.nn_model.train()")
         # nn_model = ContextUnet(n_param=1, image_size=28)
@@ -410,7 +412,7 @@ class DDPM21CM:
             # self.nn_model.load_state_dict(torch.load(config.resume)['unet_state_dict'])
             # print(f"resumed nn_model from {config.resume}")
             self.nn_model.module.load_state_dict(torch.load(config.resume)['unet_state_dict'])
-            self.nn_model.module.to(config.dtype)
+            #self.nn_model.module.to(config.dtype)
             print(f" {config.run_name} {socket.gethostbyname(socket.gethostname())} cuda:{torch.cuda.current_device()}/{self.config.global_rank} resumed nn_model from {config.resume} with {sum(x.numel() for x in self.nn_model.parameters())} parameters ".center(120,'+'))
         else:
             print(f" {config.run_name} {socket.gethostbyname(socket.gethostname())} cuda:{torch.cuda.current_device()}/{self.config.global_rank} initialized nn_model randomly with {sum(x.numel() for x in self.nn_model.parameters())} parameters ".center(120,'+'))
@@ -422,7 +424,7 @@ class DDPM21CM:
         if config.ema:
             self.ema = EMA(config.ema_rate)
             if config.resume and os.path.exists(config.resume):
-                self.ema_model = ContextUnet(n_param=config.n_param, image_size=config.HII_DIM, dim=config.dim, stride=config.stride, dtype=config.dtype).to(config.device)
+                self.ema_model = ContextUnet(n_param=config.n_param, image_size=config.HII_DIM, dim=config.dim, stride=config.stride).to(config.device)#, dtype=config.dtype
                 self.ema_model.load_state_dict(torch.load(config.resume)['ema_unet_state_dict'])
                 print(f"resumed ema_model from {config.resume}")
             else:
@@ -435,6 +437,7 @@ class DDPM21CM:
                 )
 
         self.ranges_dict = config.ranges_dict
+        self.scaler = GradScaler()
 
     def load(self):
         # rank = torch.cuda.current_device()
@@ -553,27 +556,37 @@ class DDPM21CM:
 
                 # print(f"cuda:{torch.cuda.current_device()}, x[:,0,:2,0,0] =", x[:,0,:2,0,0])
                 #with self.accelerator.accumulate(self.nn_model):
-                x = x.to(self.config.device)
+                x = x.to(self.config.device)#.to(self.config.dtype)
                 # print("x = x.to(self.config.device), x.dtype =", x.dtype)
-                x = x.to(self.config.dtype)
                 # print("x = x.to(self.dtype), x.dtype =", x.dtype)
                 # print(f"ddpm.add_noise(x), x.dtype = {x.dtype}") 
-                xt, noise, ts = self.ddpm.add_noise(x)
                 # print(f"ddpm.add_noise(x), xt.dtype = {xt.dtype}") 
-                if self.config.guide_w == -1:
-                    noise_pred = self.nn_model(xt, ts).to(x.dtype)
-                else:
-                    c = c.to(self.config.device)
-                    noise_pred = self.nn_model(xt, ts, c).to(x.dtype)
+                
+                # autocast forward propogation
+                with autocast():
+                    xt, noise, ts = self.ddpm.add_noise(x)
+
+                    if self.config.guide_w == -1:
+                        noise_pred = self.nn_model(xt, ts)#.to(x.dtype)
+                    else:
+                        c = c.to(self.config.device)
+                        noise_pred = self.nn_model(xt, ts, c)#.to(x.dtype)
                     
-                loss = F.mse_loss(noise, noise_pred)
-                loss = loss / self.config.gradient_accumulation_steps
-                loss.backward()
+                    loss = F.mse_loss(noise, noise_pred)
+                    loss = loss / self.config.gradient_accumulation_steps
+                
+                # scaler backward propogation
+                self.scaler.scale(loss).backward()
+                #loss.backward()
 
                 if (i+1) % self.config.gradient_accumulation_steps == 0:
+                    self.scaler.unscale_(self.optimizer)
                     torch.nn.utils.clip_grad_norm_(self.nn_model.parameters(), max_norm=1.0)
-                    self.optimizer.step()
+
+                    self.scaler.step(self.optimizer)
                     self.lr_scheduler.step()
+
+                    self.scaler.update()
                     self.optimizer.zero_grad()
 
                 # ema update
@@ -826,7 +839,7 @@ if __name__ == "__main__":
         max_num_img_per_gpu = args.max_num_img_per_gpu#40#2#20
         #config = TrainConfig()
         #config.world_size = world_size
-        config.dtype = torch.float32 
+        #config.dtype = torch.float32 
         config.resume = args.resume
         #config.gradient_accumulation_steps = args.gradient_accumulation_steps
         # config.resume = f"./outputs/model_state-N30-device_count3-epoch4-172.27.149.181"

perlmutter_diffusion.sbatch

@@ -5,7 +5,7 @@
 #SBATCH -q shared #regular 
 #SBATCH -N1 
 #SBATCH --gpus-per-node=1
-#SBATCH -t 0:30:00
+#SBATCH -t 0:59:00
 #SBATCH --ntasks-per-node=1
 #SBATCH -oReport-%j 
 #SBATCH --mail-type=BEGIN,END,FAIL
@@ -42,6 +42,6 @@ srun python diffusion.py \
     --gradient_accumulation_steps 1 \
     --num_new_img_per_gpu 800 \
     --max_num_img_per_gpu 80 \
-    #--resume outputs/model-N3200-device_count1-node1-epoch99-07213338 \
+    #--resume outputs/model-N3200-device_count1-node1-epoch99-16103542 \
 
 date

tensorboard.ipynb

@@ -0,0 +1,98 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "ae45e44e-a11c-43ef-b830-c7a58a72f51e",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "import nersc_tensorboard_helper\n",
+    "%load_ext tensorboard"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "a5c088b8-5051-402f-b4ec-2b684ad5a952",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "\n",
+       "      <iframe id=\"tensorboard-frame-262245829087dd6a\" width=\"100%\" height=\"800\" frameborder=\"0\">\n",
+       "      </iframe>\n",
+       "      <script>\n",
+       "        (function() {\n",
+       "          const frame = document.getElementById(\"tensorboard-frame-262245829087dd6a\");\n",
+       "          const url = new URL(\"/\", window.location);\n",
+       "          const port = 45355;\n",
+       "          if (port) {\n",
+       "            url.port = port;\n",
+       "          }\n",
+       "          frame.src = url;\n",
+       "        })();\n",
+       "      </script>\n",
+       "    "
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "%tensorboard --logdir logs --port 0"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "2f76c0a9-2218-4073-86aa-f4f655d7642f",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<a href=\"https://jupyter.nersc.gov/user/binxia/perlmutter-login-node-base/proxy/45355/\">https://jupyter.nersc.gov/user/binxia/perlmutter-login-node-base/proxy/45355/</a>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "nersc_tensorboard_helper.tb_address()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "tensorflow-2.15.0",
+   "language": "python",
+   "name": "tensorflow-2.15.0"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.18"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}