Add speedups found in profiling

train.py

@@ -130,12 +130,19 @@ def main(job_config: JobConfig):
     # loss function to be shared by Pipeline Parallel and SPMD training
     def loss_fn(pred, labels):
         return torch.nn.functional.cross_entropy(
-            pred.flatten(0, 1).float(), labels.flatten(0, 1)
+            pred.flatten(0, 1), labels.flatten(0, 1)
         )
 
-    # TODO: compiling loss function causes CUDA errors, turning off for now
-    # if job_config.training.compile:
-    #     loss_fn = torch.compile(loss_fn)
+    def train_step(input_ids, labels):
+        pred = model(input_ids)
+        loss = loss_fn(pred, labels)
+        # pred.shape=(bs, seq_len, vocab_size)
+        # need to free to before bwd to avoid peaking memory
+        del pred
+        return loss
+
+    if job_config.training.compile:
+        train_step = torch.compile(train_step)
 
     # move sharded model to CPU/GPU and initialize weights via DTensor
     if job_config.checkpoint.create_seed_checkpoint:
@@ -249,6 +256,7 @@ def loss_fn(pred, labels):
         f"total steps {job_config.training.steps} "
         f"(warmup {job_config.training.warmup_steps})"
     )
+    batch = None
     with maybe_enable_profiling(
         job_config, global_step=train_state.step
     ) as torch_profiler, maybe_enable_memory_snapshot(
@@ -260,7 +268,8 @@ def loss_fn(pred, labels):
 
             # get batch
             data_load_start = time.perf_counter()
-            batch = next(data_iterator)
+            if batch is None:
+                batch = next(data_iterator)
             input_ids, labels = batch
             ntokens_since_last_log += labels.numel()
             data_loading_times.append(time.perf_counter() - data_load_start)
@@ -306,11 +315,7 @@ def loss_fn(pred, labels):
             else:
                 # Non-PP forward / backward
                 with train_context(optional_context_parallel_ctx):
-                    pred = model(input_ids)
-                    loss = loss_fn(pred, labels)
-                    # pred.shape=(bs, seq_len, vocab_size)
-                    # need to free to before bwd to avoid peaking memory
-                    del pred
+                    loss = train_step(input_ids, labels)
                     loss.backward()
 
             # clip gradients
@@ -329,7 +334,8 @@ def loss_fn(pred, labels):
             # calculate float8 dynamic amax/scale for all-parameter for FSDP2
             # it issues a single all-reduce for all parameters at once for better performance
             float8_handler.precompute_float8_dynamic_scale_for_fsdp(model_parts)
-
+            if train_state.step < job_config.training.steps:
+                batch = next(data_iterator)
             # log metrics
             if (
                 train_state.step == 1