[float8] move inductor config setting to float8.py (#1110)

Author: danielvegamyhre

torchtitan/components/float8.py

@@ -13,6 +13,7 @@
 # Note: Performance
 # Float8 experimental is intended to be ran under `torch.compile`` for competitive performance
 
+import torch
 import torch.nn as nn
 
 from torchtitan.config_manager import JobConfig
@@ -67,6 +68,13 @@ def __init__(self, job_config: JobConfig, parallel_dims: ParallelDims):
                 f"Float8 training active with recipe {float8_config.recipe_name}"
             )
 
+            # short-term solution for https://github.com/pytorch/pytorch/issues/150859
+            if float8_config.recipe_name == "rowwise":
+                torch._inductor.config.emulate_precision_casts = True
+                logger.debug(
+                    "Set torch._inductor.config.emulate_precision_casts to True"
+                )
+
         else:
             # Mutates the model inplace replacing instances of nn.Linear with Float8Linear
             enable_fsdp_float8_all_gather = (

torchtitan/train.py

@@ -71,11 +71,6 @@ def __init__(self, job_config: JobConfig):
         if job_config.job.print_args:
             logger.info(f"Running with args: {job_config.to_dict()}")
 
-        # short-term solution for https://github.com/pytorch/pytorch/issues/150859
-        if job_config.float8.recipe_name == "rowwise":
-            torch._inductor.config.emulate_precision_casts = True
-            logger.debug("Set torch._inductor.config.emulate_precision_casts to True")
-
         # take control of garbage collection to avoid stragglers
         self.gc_handler = utils.GarbageCollection(gc_freq=job_config.training.gc_freq)