rebase

Signed-off-by: Jonathan Mitchell <jomitchell@nvidia.com>

Author: jomitchellnv

bionemo-recipes/recipes/llama3_native_te/train_ddp.py

@@ -119,7 +119,11 @@ def main(args: DictConfig) -> float | None:
     with transformer_engine.pytorch.quantized_model_init(
         recipe=fp8_recipe, **args.fp8_config.quantized_model_init_kwargs
     ):
-        model = model_class(config)
+        model = (
+            model_class(config, fp8_recipe=fp8_recipe, fp4_recipe=fp4_recipe)
+            if model_class is NVLlamaForCausalLM
+            else model_class(config)
+        )
 
     logger.info("Initialized Model:\n%s", model)
 

bionemo-recipes/recipes/llama3_native_te/train_fsdp2.py

@@ -128,7 +128,11 @@ def main(args: DictConfig) -> float | None:
             recipe=fp8_recipe, **args.fp8_config.quantized_model_init_kwargs
         ),
     ):
-        model = model_class(config)
+        model = (
+            model_class(config, fp8_recipe=fp8_recipe, fp4_recipe=fp4_recipe)
+            if model_class is NVLlamaForCausalLM
+            else model_class(config)
+        )
 
     logger.info("Initialized Model:\n%s", model)
 

bionemo-recipes/recipes/llama3_native_te/train_fsdp2_cp.py

@@ -127,7 +127,7 @@ def main(args: DictConfig) -> float | None:
             recipe=fp8_recipe, **args.fp8_config.quantized_model_init_kwargs
         ),
     ):
-        model = NVLlamaForCausalLM(config)
+        model = NVLlamaForCausalLM(config, fp8_recipe=fp8_recipe, fp4_recipe=fp4_recipe)
 
     logger.info("Initialized Model:\n%s", model)