[mxfp8 training] add TP warning

Author: danielvegamyhre

torchtitan/components/quantization/mx.py

@@ -84,6 +84,12 @@ def __init__(
                 "torch.compile enablement is required for highest performance of MXFP8 dynamic quantization."
             )
 
+        if parallel_dims.tp > 1:
+            logger.warning(
+                "TP support for MXFP8 linears is still in progress. "
+                "Any linear layers with TP applied will use default precision, not MXFP8."
+            )
+
         self.config = config
         self.enabled = True
         logger.info("MXFP8 MoE training enabled")

torchtitan/models/deepseek_v3/config_registry.py

@@ -12,6 +12,7 @@
     Float8GroupedMMConverter,
     Float8LinearConverter,
 )
+from torchtitan.components.quantization.mx import MXFP8Converter
 from torchtitan.config import (
     ActivationCheckpointConfig,
     CompileConfig,
@@ -82,9 +83,17 @@ def deepseek_v3_16b() -> Trainer.Config:
             seq_len=4096,
             steps=1000,
         ),
+        model_converters=ModelConvertersContainer.Config(
+            converters=[
+                MXFP8Converter.Config(fqns=["moe.shared_experts"]),
+            ],
+            print_after_conversion=True,
+        ),
         parallelism=ParallelismConfig(
             pipeline_parallel_schedule="Interleaved1F1B",
-            expert_parallel_degree=8,
+            expert_parallel_degree=4,
+            data_parallel_shard_degree=4,
+            tensor_parallel_degree=2,
             expert_tensor_parallel_degree=1,
         ),
         checkpoint=CheckpointManager.Config(interval=10),