diff --git a/torchtitan/experiments/torchcomms/integration_tests.py b/torchtitan/experiments/torchcomms/integration_tests.py index c16778d5a2..5df527a064 100755 --- a/torchtitan/experiments/torchcomms/integration_tests.py +++ b/torchtitan/experiments/torchcomms/integration_tests.py @@ -30,10 +30,11 @@ def build_torchcomms_test_list() -> list[OverrideDefinitions]: [ "--parallelism.tensor_parallel_degree 2", "--parallelism.pipeline_parallel_degree 2", + "--compile.enable", ], ], - "FSDP+TP+PP", - "3d_dp+tp+pp", + "FSDP+TP+PP+compile", + "3d_dp+tp+pp+compile", ngpu=8, ), OverrideDefinitions( @@ -41,10 +42,11 @@ def build_torchcomms_test_list() -> list[OverrideDefinitions]: [ "--parallelism.context_parallel_degree 2", "--parallelism.pipeline_parallel_degree 2", + "--compile.enable", ], ], - "DP+CP+PP", - "3d_dp+cp+pp", + "DP+CP+PP+compile", + "3d_dp+cp+pp+compile", ngpu=8, ), # TODO: Enable async TP test once fixes are available for running on CI nodes.