wip

app.py

@@ -13,17 +13,19 @@ from gated_state_spaces_pytorch import GatedStateSpacesLM
 from gated_state_spaces_pytorch.autoregressive_wrapper import AutoregressiveWrapper
 
 from c4x import C4X
+from accelerate import Accelerator
 
 
-if __name__ == '__main__':
-    wandb.init(
-        project="gated-state-space", 
-        entity="naxalpha",
+def main():
+    accelerator = Accelerator(
+        gradient_accumulation_steps=4,
     )
     
-    # gpt_2 = GPT2LMHeadModel.from_pretrained('gpt2-xl')
-    # gpt_2.requires_grad_(False)
-    # gpt_2 = gpt_2.cuda()
+    if accelerator.is_main_process:
+        wandb.init(
+            project="gated-state-space", 
+            entity="naxalpha",
+        )
 
     f_emb = 1600
     model = AutoregressiveWrapper(
@@ -32,56 +34,57 @@ if __name__ == '__main__':
             dim=f_emb,
             depth=24,
         ),
-    )
-    wandb.watch(model)
-
-    # emb = gpt_2.state_dict()['transformer.wte.weight']
-
+    )        
     model.net.token_emb.weight.requires_grad_(False)
-    # model.net.token_emb.weight.copy_(emb)
-
     model.net.to_logits.weight.requires_grad_(False)
-    # model.net.to_logits.weight.copy_(emb)
-    
     model.net.to_logits = nn.Sequential(
         nn.LayerNorm(f_emb),
         model.net.to_logits,
     )
+    model = model.to(accelerator.device)
+    
+    if accelerator.is_main_process:
+        wandb.watch(model)
     
     model.load_state_dict(torch.load('model.pt'))
-    model = model.cuda()
     optim = AdamW(model.parameters(), 2e-5)
 
-    bs = 8
+    bs = 16
     kk = 128
     dsx = C4X(kk+1)
     dlx = DataLoader(
         dsx,
         batch_size=bs, 
-        num_workers=16,
+        num_workers=8,
     )
 
     k = 4
-    prog = tqdm(dlx)
+    prog = tqdm(dlx, disable=not accelerator.is_main_process)
+    
+    model, optim, dlx = accelerator.prepare(model, optim, dlx)
+    
     optim.zero_grad()
-
     for i, batch in enumerate(prog):
-        batch = batch.cuda()
-        los = model(batch)
-
-        (los / k).backward()
-        if (i+1) % k == 0:
-            clip_grad_norm_(
-                model.parameters(),
-                max_norm=1.,
-            )
+        batch = batch.to(accelerator.device)
+        with accelerator.accumulate(model):
+            with accelerator.autocast():
+                los = model(batch)
+            accelerator.backward(los)
+            if accelerator.sync_gradients:
+                accelerator.clip_grad_norm_(
+                    model.parameters(), 
+                    1.0,
+                )
             optim.step()
             optim.zero_grad()
 
-        if i % 1000 == 0:
+        if i % 1000 == 0 and accelerator.is_main_process:
+            print('generating...')
+            accelerator.wait_for_everyone()
+            unwrapped_model = accelerator.unwrap_model(model)
             b, n = 4, 512
-            init = torch.tensor([[50256]]*b).cuda()
-            prd = model.generate(init, n)
+            init = torch.tensor([[50256]]*b).to(accelerator.device)
+            prd = unwrapped_model.generate(init, n)
             prd = [dsx.decode(p) for p in prd]
             try:
                 wandb.log(dict(
@@ -92,9 +95,14 @@ if __name__ == '__main__':
                 )), step=i)
             except Exception as ex:
                 print('Failed to log to W&B...', ex)
-            torch.save(model.state_dict(), 'model.pt')
+            accelerator.save(unwrapped_model.state_dict(), 'model.pt')
+        
+        if i % 10 == 0 and accelerator.is_main_process:
+            print('logging...')
+            wandb.log(dict(
+                loss=los.item(),
+            ), step=i)
+            prog.set_postfix(loss=los.item())
 
-        wandb.log(dict(
-            loss=los.item(),
-        ), step=i)
-        prog.set_postfix(loss=los.item())
+if __name__ == '__main__':
+    main()