hit ddp training scripts as well

Signed-off-by: Peter St. John <pstjohn@nvidia.com>

Author: pstjohn

bionemo-recipes/recipes/esm2_native_te/perf_logger.py

@@ -143,3 +143,6 @@ def finish(self):
 
         wandb.finish()
         self._progress_bar.close()
+
+        if self.fp8_stats_enabled:
+            debug_api.end_debug()

bionemo-recipes/recipes/esm2_native_te/train_ddp.py

@@ -30,6 +30,7 @@
 from checkpoint import load_checkpoint_ddp, save_checkpoint_ddp, save_final_model_ddp, should_save_checkpoint
 from dataset import create_bshd_dataloader, create_thd_dataloader
 from distributed_config import DistributedConfig
+from fp8_debugging import initialize_fp8_debugging
 from perf_logger import PerfLogger
 from scheduler import get_linear_schedule_with_warmup
 
@@ -52,24 +53,10 @@ def main(args: DictConfig) -> float | None:
     torch.distributed.init_process_group(backend="nccl", device_id=device)
     torch.cuda.set_device(dist_config.local_rank)
 
-    # TE Debug feature logging
-    if args.fp8_stats_config.enabled and not args.fp8_config.enabled:
-        raise ValueError(
-            "fp8_stats_config.enabled is true but fp8_config.enabled is false, please set fp8_config.enabled to true in the config if you wish to collect FP8 stats"
-        )
-
+    # TE Debug feature logging - MUST be done BEFORE FSDP wrapping
     if args.fp8_stats_config.enabled:
-        fp8_stats_file = args.fp8_stats_config.fp8_stats_file
-        fp8_log_dir = Path(args.fp8_stats_config.fp8_log_dir) / f"rank_{dist_config.rank}"
-        fp8_log_dir.mkdir(parents=True, exist_ok=True)
-        logger.info(f"Logging FP8 stats to {fp8_log_dir}")
-        te_features_dir = str(Path(transformer_engine.__file__).parent / "debug" / "features")
-        debug_api.initialize(
-            config_file=fp8_stats_file,
-            feature_dirs=[te_features_dir],
-            log_dir=fp8_log_dir,
-            default_logging_enabled=True,
-        )
+        initialize_fp8_debugging(dist_config, **args.fp8_stats_config, fp8_enabled=args.fp8_config.enabled)
+
     # Create a device mesh for DDP. While this isn't strictly necessary, it mirrors the device mesh we create for FSDP2
     # and MFSDP.
     device_mesh = init_device_mesh("cuda", mesh_shape=(dist_config.world_size,), mesh_dim_names=("ddp",))
@@ -157,8 +144,6 @@ def main(args: DictConfig) -> float | None:
             loss = outputs.loss
             loss.backward()
 
-            if args.fp8_stats_config.enabled:
-                debug_api.step()
             # Compute and clip gradient norms.
             total_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0).item()
 
@@ -206,8 +191,6 @@ def main(args: DictConfig) -> float | None:
 
     # Clean up distributed training
     perf_logger.finish()
-    if args.fp8_stats_config.enabled:
-        debug_api.end_debug()
     torch.distributed.destroy_process_group()
 
     return perf_logger.min_loss

bionemo-recipes/recipes/esm2_native_te/train_fsdp2.py

@@ -205,8 +205,6 @@ def main(args: DictConfig) -> float | None:
 
     # Clean up distributed training
     perf_logger.finish()
-    if args.fp8_stats_config.enabled:
-        debug_api.end_debug()
     torch.distributed.destroy_process_group()
 
     return perf_logger.min_loss

bionemo-recipes/recipes/llama3_native_te/perf_logger.py

@@ -177,6 +177,9 @@ def finish(self):
         wandb.finish()
         self._progress_bar.close()
 
+        if self.fp8_stats_enabled:
+            debug_api.end_debug()
+
 
 def setup_profiler(args: DictConfig, wandb_run: wandb.Run):
     """Setup a basic torch profiler for the experiment.

bionemo-recipes/recipes/llama3_native_te/train_ddp.py

@@ -32,6 +32,7 @@
 from checkpoint import load_checkpoint_ddp, save_checkpoint_ddp, save_final_model_ddp, should_save_checkpoint
 from dataset import create_bshd_dataloader, create_thd_dataloader
 from distributed_config import DistributedConfig
+from fp8_debugging import initialize_fp8_debugging
 from modeling_llama_te import NVLlamaConfig, NVLlamaForCausalLM
 from perf_logger import PerfLogger
 from scheduler import get_cosine_annealing_schedule_with_warmup
@@ -56,23 +57,8 @@ def main(args: DictConfig) -> float | None:
     torch.cuda.set_device(dist_config.local_rank)
 
     # TE Debug feature logging
-    if args.fp8_stats_config.enabled and not args.fp8_config.enabled:
-        raise ValueError(
-            "fp8_stats_config.enabled is true but fp8_config.enabled is false, please set fp8_config.enabled to true in the config if you wish to collect FP8 stats"
-        )
-
     if args.fp8_stats_config.enabled:
-        fp8_stats_file = args.fp8_stats_config.fp8_stats_file
-        fp8_log_dir = Path(args.fp8_stats_config.fp8_log_dir) / f"rank_{dist_config.rank}"
-        fp8_log_dir.mkdir(parents=True, exist_ok=True)
-        logger.info(f"Logging FP8 stats to {fp8_log_dir}")
-        te_features_dir = str(Path(transformer_engine.__file__).parent / "debug" / "features")
-        debug_api.initialize(
-            config_file=fp8_stats_file,
-            feature_dirs=[te_features_dir],
-            log_dir=fp8_log_dir,
-            default_logging_enabled=True,
-        )
+        initialize_fp8_debugging(dist_config, **args.fp8_stats_config, fp8_enabled=args.fp8_config.enabled)
 
     # Create a device mesh for DDP. While this isn't strictly necessary, it mirrors the device mesh we create for FSDP2
     # and MFSDP.
@@ -163,8 +149,6 @@ def main(args: DictConfig) -> float | None:
 
                 # Log microbatch step data for accumulation metrics
                 perf_logger.log_micro_step(batch=batch, outputs=outputs)
-            if args.fp8_stats_config.enabled:
-                debug_api.step()
 
             # Gradient accumulation - only step optimizer after accumulating gradients
             if micro_step % args.grad_acc_steps == 0:
@@ -215,8 +199,6 @@ def main(args: DictConfig) -> float | None:
 
     # Clean up distributed training
     perf_logger.finish()
-    if args.fp8_stats_config.enabled:
-        debug_api.end_debug()
     torch.distributed.destroy_process_group()
 
     return perf_logger.min_loss

bionemo-recipes/recipes/llama3_native_te/train_fsdp2.py

@@ -229,8 +229,6 @@ def main(args: DictConfig) -> float | None:
 
     # Clean up distributed training
     perf_logger.finish()
-    if args.fp8_stats_config.enabled:
-        debug_api.end_debug()
     torch.distributed.destroy_process_group()
 
     return perf_logger.min_loss