add tensorboard to training script

Summary:
- add tensorboard integration and separate the metrics by run id and replica id
- have an output folder per replica id

Author: tushar00jain

train_diloco.py

@@ -24,6 +24,7 @@
 from torch.distributed.pipelining import SplitPoint, pipeline
 from torch.export import export
 from torchdata.stateful_dataloader import StatefulDataLoader
+from torch.utils.tensorboard import SummaryWriter
 
 from torchft import (
     DistributedSampler,
@@ -41,7 +42,11 @@
 @record
 def main() -> None:
     REPLICA_GROUP_ID = int(os.environ.get("REPLICA_GROUP_ID", 0))
-    NUM_REPLICA_GROUPS = int(os.environ.get("NUM_REPLICA_GROUPS", 2))
+    RUN = int(os.environ.get("RUN", 0))
+
+    output_folder = f"output/replica-{REPLICA_GROUP_ID}"
+
+    writer = SummaryWriter(f"{output_folder}/tensorboard", max_queue=1000)
 
     def load_state_dict(state_dict):
         m.load_state_dict(state_dict["model"])
@@ -175,7 +180,7 @@ def forward(self, x):
 
     def trace_handler(p):
         p.export_chrome_trace(
-            f"/home/tushar00jain/trace_{p.step_num}_{REPLICA_GROUP_ID}.json"
+            f"{output_folder}/profiles/step-{p.step_num}.json"
         )
 
     # You can use an epoch based training but with faults it's easier to use step
@@ -188,6 +193,7 @@ def trace_handler(p):
     )
 
     prof.start()
+    tensorboard_key_prefix = f"Run:{RUN}"
     with DiLoCo(
         manager,
         module_partitions if USE_STREAMING else [m],
@@ -210,16 +216,21 @@ def trace_handler(p):
                 out = m(inputs)
                 loss = criterion(out, labels)
 
+                writer.add_scalar(f"{tensorboard_key_prefix}/loss", loss, i)
+
                 loss.backward()
 
                 inner_optimizer.step()
 
+                writer.add_scalar(f"{tensorboard_key_prefix}/num_participants", manager.num_participants(), i)
+                writer.add_scalar(f"{tensorboard_key_prefix}/current_step", manager.current_step(), i)
                 if manager.current_step() % 100 == 0:
                     print(f"[{manager.current_step()}] loss = {loss.item()}")
 
                 if manager.current_step() >= 15:
                     # complete training
                     prof.stop()
+                    writer.flush()
                     exit()