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"