Skip to content

Fix ci hang in torch2.7& improve ut #7321

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
Jun 2, 2025
Merged

Fix ci hang in torch2.7& improve ut #7321

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
647c2a5
format
inkcherry May 30, 2025
afd3cf5
update
inkcherry May 30, 2025
5374485
clean up
inkcherry May 30, 2025
3db9c27
clean up
inkcherry May 30, 2025
3320fbd
Merge branch 'master' into fix_cihang
sfc-gh-truwase May 31, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 11 additions & 7 deletions tests/unit/model_parallelism/test_autotp_training.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,9 @@ def should_assert_with_msg(expected_message):
pass
else:
raise e
else:
raise AssertionError(f"Expected AssertionError with message '{expected_message}' "
"but no exception was raised.")


@pytest.mark.parametrize("tp_size", [2, 4])
Expand All @@ -67,7 +70,6 @@ class TestTpParallelStates(DistributedTest):
def test(self, tp_size: int):
skip_on_device()
set_autotp_mode(training=True)

dp_size = 4 / tp_size
hidden_dim = 128
config_dict = {
Expand All @@ -88,7 +90,7 @@ def test(self, tp_size: int):
@pytest.mark.parametrize("tp_size", [2, 4])
class TestTpDataloaderCorrectness(DistributedTest):
world_size = 4
reuse_dist_env = True
reuse_dist_env = False

def test(self, tp_size: int):
skip_on_device()
Expand Down Expand Up @@ -117,6 +119,8 @@ def test(self, tp_size: int):

model = SimpleModel(hidden_dim=hidden_dim)
model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
torch.manual_seed(42)

data_loader = random_dataloader(model=model,
total_samples=3,
hidden_dim=hidden_dim,
Expand Down Expand Up @@ -165,7 +169,7 @@ def process_linear_layer(hidden_dim, input):
@pytest.mark.parametrize("tp_overlap_comm", [True, False])
class TestTpLayerFwdBwd(DistributedTest):
world_size = 4
reuse_dist_env = True
reuse_dist_env = False

def testRowParallel(self, tp_size: int, tp_overlap_comm: bool):
skip_on_device()
Expand Down Expand Up @@ -278,10 +282,10 @@ def testColumnParallel(self, tp_size: int, tp_overlap_comm: bool):
atol=1e-2)


@pytest.mark.sequential
# @pytest.mark.sequential
class TestParamsGather(DistributedTest):
world_size = 4
reuse_dist_env = True
reuse_dist_env = False

@pytest.mark.parametrize("layer_type", ["linear", "linearallreduce"])
def test(self, layer_type):
Expand Down Expand Up @@ -388,7 +392,7 @@ def prepare_tp_model(hidden_dim, nlayers, linear_indices, allreduce_indices, gro
class TestSave(DistributedTest):

world_size = 4
reuse_dist_env = True
reuse_dist_env = False

def test_save_original_weight(self, tp_size: int, zero_stage: int):
skip_on_device()
Expand Down Expand Up @@ -520,7 +524,7 @@ def test_ckpt_save(self, tmpdir, tp_size: int, zero_stage: int):
class TestTpGradNorm(DistributedTest):

world_size = 4
reuse_dist_env = True
reuse_dist_env = False

def test(self, tp_size: int, zero_stage: int):
skip_on_device()
Expand Down
Loading