Remove HSDP validation check

[mvpatel2000]

Currently, HSDP validates that all intra/inter node PGs are the same. This makes sense if you are only using HSDP with no other forms of parallelism and is a nice but not necessary sanity check.

However, if you want to mix HSDP with other forms, say tensor parallelism on the FFN of a transformer block, the intra/inter node PGs will be different for that layer. This check raises errors in this scenario, so we need to remove this assumption.

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu @fegin @XilunWu @wanchaol @fduwjj @wz337 @tianyu-l @wconstab @yf225

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/112435

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit e08159e with merge base d444a3b ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[awgu]

If I understand correctly, there is still some value in the check; however, it is currently overly strict and is problematic for manual wrapping + HSDP. I think there was someone internally working on relaxing the check.

The check that is valuable is that if you are using HSDP, then each HSDP instance should use the same process groups if using the same ranks. We do not want to create a different pair of process groups per HSDP instance.

[mvpatel2000]

@awgu got it! If someone is working on it, feel free to close this PR then :)

[awgu]

Let me follow-up on the progress on that PR and get back to you!

[awgu]

@fegin @wz337 Is there anything from the checkpointing side that requires each FSDP instance to use the HSDP process groups?

If not, then I think removing this requirement/check sounds good to me (and we would need to remove the unit test).

pytorch/test/distributed/fsdp/test_fsdp_hybrid_shard.py

Lines 119 to 157 in b8b3c26

	@skip_if_lt_x_gpu(2)
	def test_hybrid_shard_pg_mismatch_raises(self):
	model = MyModel().cuda()
	intra_pg = self.process_group
	inter_pg = dist.new_group(ranks=[self.rank])
	# Mismatched process groups for intra-node
	model.lin1 = FSDP(
	model.lin1,
	process_group=(intra_pg, inter_pg),
	sharding_strategy=ShardingStrategy.HYBRID_SHARD,
	)
	model = FSDP(
	model,
	process_group=(dist.new_group(), dist.new_group()),
	sharding_strategy=ShardingStrategy.HYBRID_SHARD,
	)
	# Errors during _lazy_init
	inp = torch.randn(4, 10)
	with self.assertRaisesRegex(
	ValueError, "intra-node process groups do not match"
	):
	model(inp)
	# Mismatched process groups for inter-node
	model = MyModel().cuda()
	model.lin1 = FSDP(
	model.lin1,
	process_group=(intra_pg, inter_pg),
	sharding_strategy=ShardingStrategy.HYBRID_SHARD,
	)
	model = FSDP(
	model,
	process_group=(intra_pg, dist.new_group()),
	sharding_strategy=ShardingStrategy.HYBRID_SHARD,
	)
	with self.assertRaisesRegex(
	ValueError, "inter-node process groups do not match"
	):
	model(inp)

[mvpatel2000]

@fegin @wz337 Is there anything from the checkpointing side that requires each FSDP instance to use the HSDP process groups?

If not, then I think removing this requirement/check sounds good to me (and we would need to remove the unit test).

pytorch/test/distributed/fsdp/test_fsdp_hybrid_shard.py

Lines 119 to 157 in b8b3c26

	@skip_if_lt_x_gpu(2)
	def test_hybrid_shard_pg_mismatch_raises(self):
	model = MyModel().cuda()
	intra_pg = self.process_group
	inter_pg = dist.new_group(ranks=[self.rank])
	# Mismatched process groups for intra-node
	model.lin1 = FSDP(
	model.lin1,
	process_group=(intra_pg, inter_pg),
	sharding_strategy=ShardingStrategy.HYBRID_SHARD,
	)
	model = FSDP(
	model,
	process_group=(dist.new_group(), dist.new_group()),
	sharding_strategy=ShardingStrategy.HYBRID_SHARD,
	)
	# Errors during _lazy_init
	inp = torch.randn(4, 10)
	with self.assertRaisesRegex(
	ValueError, "intra-node process groups do not match"
	):
	model(inp)
	# Mismatched process groups for inter-node
	model = MyModel().cuda()
	model.lin1 = FSDP(
	model.lin1,
	process_group=(intra_pg, inter_pg),
	sharding_strategy=ShardingStrategy.HYBRID_SHARD,
	)
	model = FSDP(
	model,
	process_group=(intra_pg, dist.new_group()),
	sharding_strategy=ShardingStrategy.HYBRID_SHARD,
	)
	with self.assertRaisesRegex(
	ValueError, "inter-node process groups do not match"
	):
	model(inp)

@awgu @fegin @wz337 bumping this request! would love to have this issue resolved

[mvpatel2000]

@awgu @fegin @wz337 bumping this please!

[pytorch-bot]

Please seek CI approval before scheduling CIFlow labels

[fegin]

I think it is okay to remove the check. Will let @wz337 to review again.

[wz337]

@fegin @wz337 Is there anything from the checkpointing side that requires each FSDP instance to use the HSDP process groups?

If not, then I think removing this requirement/check sounds good to me (and we would need to remove the unit test).

pytorch/test/distributed/fsdp/test_fsdp_hybrid_shard.py

Lines 119 to 157 in b8b3c26

	@skip_if_lt_x_gpu(2)
	def test_hybrid_shard_pg_mismatch_raises(self):
	model = MyModel().cuda()
	intra_pg = self.process_group
	inter_pg = dist.new_group(ranks=[self.rank])
	# Mismatched process groups for intra-node
	model.lin1 = FSDP(
	model.lin1,
	process_group=(intra_pg, inter_pg),
	sharding_strategy=ShardingStrategy.HYBRID_SHARD,
	)
	model = FSDP(
	model,
	process_group=(dist.new_group(), dist.new_group()),
	sharding_strategy=ShardingStrategy.HYBRID_SHARD,
	)
	# Errors during _lazy_init
	inp = torch.randn(4, 10)
	with self.assertRaisesRegex(
	ValueError, "intra-node process groups do not match"
	):
	model(inp)
	# Mismatched process groups for inter-node
	model = MyModel().cuda()
	model.lin1 = FSDP(
	model.lin1,
	process_group=(intra_pg, inter_pg),
	sharding_strategy=ShardingStrategy.HYBRID_SHARD,
	)
	model = FSDP(
	model,
	process_group=(intra_pg, dist.new_group()),
	sharding_strategy=ShardingStrategy.HYBRID_SHARD,
	)
	with self.assertRaisesRegex(
	ValueError, "inter-node process groups do not match"
	):
	model(inp)

We are relying on the DTensor to do all_gather and chunk so we don't use the HSDP process groups directly. So I think it should be fine removing this requirement.

[wz337]

I think it's ok to remove this check.

Could you also include removing the unit test in the PR as @awgu mentioned so CI doesn't break?

https://github.com/pytorch/pytorch/blob/main/test/distributed/fsdp/test_fsdp_hybrid_shard.py#L120

[pytorch-bot]

Please seek CI approval before scheduling CIFlow labels

[mvpatel2000]

@wz337 test removed!

@wconstab sorry -- updated to be more clear :)

[mvpatel2000]

@pytorchmergebot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

[pytorchmergebot]

Merge failed

Reason: 3 mandatory check(s) failed. The first few are:

• BC Lint
• pull
• Lint

Dig deeper by viewing the failures on hud

Details for Dev Infra team

Raised by workflow job

Failing merge rule: Core Maintainers

[Skylion007]

@pytorchbot merge -r

[pytorchmergebot]

@pytorchbot started a rebase job onto refs/remotes/origin/viable/strict. Check the current status here

[pytorchmergebot]

Successfully rebased mvpatel2000/remove-hsdp-validate onto refs/remotes/origin/viable/strict, please pull locally before adding more changes (for example, via git checkout mvpatel2000/remove-hsdp-validate && git pull --rebase)

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here