From bb94c37eefbc384c3418ad6364f2fa4b3bb0979b Mon Sep 17 00:00:00 2001 From: Daniel Vega-Myhre Date: Tue, 8 Jul 2025 15:44:49 -0700 Subject: [PATCH] support TP --- torchtitan/components/quantization/float8.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/torchtitan/components/quantization/float8.py b/torchtitan/components/quantization/float8.py index 0d8d4460d..a5b2f967c 100644 --- a/torchtitan/components/quantization/float8.py +++ b/torchtitan/components/quantization/float8.py @@ -59,9 +59,6 @@ def __init__(self, job_config: JobConfig, parallel_dims: ParallelDims): # Validate MoE training prototype limitations. if self.moe_fqns: - assert ( - job_config.parallelism.tensor_parallel_degree == 1 - ), "Float8 MoE training prototype does not yet support tensor parallelism" assert ( job_config.parallelism.pipeline_parallel_degree == 1 ), "Float8 MoE training prototype does not yet support pipeline parallelism"