Add BSHD packed dataloader toggle and FP8 test

- Update train_fsdp2.py and train_ddp.py to toggle between dataloaders:
  - use_sequence_packing=true + attn_input_format=bshd -> BSHD packed
  - use_sequence_packing=true + attn_input_format=thd -> THD packed
  - use_sequence_packing=false -> BSHD unpacked
- Add test_train_fsdp2_fp8_bshd_packed test for FP8 with BSHD packing

Signed-off-by: Savitha Srinivasan <savithas@nvidia.com>

Author: savitha-eng

bionemo-recipes/recipes/llama3_native_te/tests/test_train.py

@@ -431,6 +431,28 @@ def test_train_fsdp2_fp8_first_last_bf16(tmp_path, recipe_path):
     assert final_loss < 8.0, f"Final loss {final_loss} is too high, expected < 8.0"
 
 
+def test_train_fsdp2_fp8_bshd_packed(tmp_path, recipe_path):
+    """Test that FSDP2 training works with FP8 enabled and BSHD packed dataloader."""
+    with initialize_config_dir(config_dir=str(recipe_path / "hydra_config"), version_base="1.2"):
+        sanity_config = compose(
+            config_name="L0_sanity",
+            overrides=[
+                f"+wandb.dir={tmp_path}",
+                f"checkpoint.ckpt_dir={tmp_path}",
+                "fp8_config.enabled=true",
+                "use_sequence_packing=true",
+                "config_kwargs.attn_input_format=bshd",
+                "+dataset.pad_to_multiple_of=16",
+            ],
+        )
+
+    final_loss = main_fsdp2(sanity_config)
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    assert final_loss < 8.0, f"Final loss {final_loss} is too high, expected < 8.0"
+
+
 @requires_datacenter_hardware
 def test_sanity_fsdp2_cp(tmp_path, recipe_path):
     # Run the training script with Hydra configuration overrides

bionemo-recipes/recipes/llama3_native_te/train_ddp.py

@@ -28,7 +28,7 @@
 from transformers.models.llama.modeling_llama import LlamaForCausalLM
 
 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 dataset import create_bshd_dataloader, create_bshd_packed_dataloader, create_thd_dataloader
 from distributed_config import DistributedConfig
 from modeling_llama_te import NVLlamaConfig, NVLlamaForCausalLM
 from perf_logger import PerfLogger
@@ -93,8 +93,14 @@ def main(args: DictConfig) -> float | None:
     )
 
     if args.use_sequence_packing:
-        train_dataloader, dataset_or_sampler = create_thd_dataloader(dist_config, **args.dataset)
+        if args.config_kwargs.attn_input_format == "bshd":
+            # BSHD with full packing (cross-boundary attention, no cu_seqlens)
+            train_dataloader, dataset_or_sampler = create_bshd_packed_dataloader(dist_config, **args.dataset)
+        else:
+            # THD with packing (respects boundaries via cu_seqlens)
+            train_dataloader, dataset_or_sampler = create_thd_dataloader(dist_config, **args.dataset)
     else:
+        # Standard BSHD with windowing (no packing)
         train_dataloader, dataset_or_sampler = create_bshd_dataloader(dist_config, **args.dataset)
 
     if args.use_torch_compile:

bionemo-recipes/recipes/llama3_native_te/train_fsdp2.py

@@ -36,7 +36,7 @@
     save_final_model_fsdp2,
     should_save_checkpoint,
 )
-from dataset import create_bshd_dataloader, create_thd_dataloader
+from dataset import create_bshd_dataloader, create_bshd_packed_dataloader, create_thd_dataloader
 from distributed_config import DistributedConfig
 from modeling_llama_te import NVLlamaConfig, NVLlamaForCausalLM
 from perf_logger import PerfLogger
@@ -110,8 +110,14 @@ def main(args: DictConfig) -> float | None:
     scheduler = get_cosine_annealing_schedule_with_warmup(optimizer, **args.lr_scheduler_kwargs)
 
     if args.use_sequence_packing:
-        train_dataloader, dataset_or_sampler = create_thd_dataloader(dist_config, **args.dataset)
+        if args.config_kwargs.attn_input_format == "bshd":
+            # BSHD with full packing (cross-boundary attention, no cu_seqlens)
+            train_dataloader, dataset_or_sampler = create_bshd_packed_dataloader(dist_config, **args.dataset)
+        else:
+            # THD with packing (respects boundaries via cu_seqlens)
+            train_dataloader, dataset_or_sampler = create_thd_dataloader(dist_config, **args.dataset)
     else:
+        # Standard BSHD with windowing (no packing)
         train_dataloader, dataset_or_sampler = create_bshd_dataloader(dist_config, **args.dataset)
 
     if args.use_torch_compile: