distributed training

app.py

@@ -18,14 +18,10 @@ from accelerate import Accelerator
 
 def main():
     accelerator = Accelerator(
+        log_with="wandb",
         gradient_accumulation_steps=4,
     )
-    
-    if accelerator.is_main_process:
-        wandb.init(
-            project="gated-state-space", 
-            entity="naxalpha",
-        )
+    accelerator.init_trackers("gated-state-space")
 
     f_emb = 1600
     model = AutoregressiveWrapper(
@@ -49,7 +45,7 @@ def main():
     model.load_state_dict(torch.load('model.pt'))
     optim = AdamW(model.parameters(), 2e-5)
 
-    bs = 16
+    bs = 24
     kk = 128
     dsx = C4X(kk+1)
     dlx = DataLoader(
@@ -58,7 +54,6 @@ def main():
         num_workers=8,
     )
 
-    k = 4
     prog = tqdm(dlx, disable=not accelerator.is_main_process)
     
     model, optim, dlx = accelerator.prepare(model, optim, dlx)
@@ -78,8 +73,7 @@ def main():
             optim.step()
             optim.zero_grad()
 
-        if i % 1000 == 0 and accelerator.is_main_process:
-            print('generating...')
+        if i % 1000 == 0:
             accelerator.wait_for_everyone()
             unwrapped_model = accelerator.unwrap_model(model)
             b, n = 4, 512
@@ -87,19 +81,18 @@ def main():
             prd = unwrapped_model.generate(init, n)
             prd = [dsx.decode(p) for p in prd]
             try:
-                wandb.log(dict(
+                accelerator.log(dict(
                     text=wandb.Html(
                         '<hr>'.join(
                             p.replace('\n', '<br>') for p in prd
                         )
                 )), step=i)
             except Exception as ex:
-                print('Failed to log to W&B...', ex)
-            accelerator.save(unwrapped_model.state_dict(), 'model.pt')
+                accelerator.print('Failed to log to W&B...', ex)
+            accelerator.save(unwrapped_model.state_dict(), 'model2.pt')
         
-        if i % 10 == 0 and accelerator.is_main_process:
-            print('logging...')
-            wandb.log(dict(
+        if i % 10 == 0:
+            accelerator.log(dict(
                 loss=los.item(),
             ), step=i)
             prog.set_postfix(loss=los.item())