[FSDP][Easy] Rename streams; add back stream sharing test

[awgu]

Stack from ghstack (oldest at bottom):

• [FSDP][9/N] Introduce CustomPolicy #104986
• [FSDP][8/N] Replace _FSDPPolicy.policy with _Policy._run_policy #104969
• [FSDP][Easy] Allow ModuleWrapPolicy to take Iterable #104999
• [FSDP][7/N] Add warning about frozen params #104967
• [FSDP][6/N] Check valid param freezing for ModuleWrapPolicy #104427
• -> [FSDP][Easy] Rename streams; add back stream sharing test #104966

Purely out of preference, this PR renames the streams to _unshard_stream instead of _streams_unshard etc. since the former reads more naturally. The PR also removes some duplicated comments and adds back a unit test that streams are shared.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 b57126f:
💚 Looks good so far! There are no failures yet. 💚

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

[awgu]

We do not need to duplicate these comments with the initial construction in _init_streams(), as the comments may become stale and here we are only sharing the state.

pytorch/torch/distributed/fsdp/_runtime_utils.py

Lines 305 to 324 in b57126f

	def _init_streams(
	state: _FSDPState,
	) -> _FSDPState:
	"""
	Initializes CUDA streams for overlapping communication, computation, and
	data transfers. The streams should be shared across FSDP instances.
	"""
	assert state._is_root
	assert state._device_handle.is_available()
	# Stream for unshard logic, including allocating the all-gather destination
	# tensors and the all-gathers themselves.
	state._unshard_stream = state._device_handle.Stream()
	# Stream for overlapping gradient reduction with the backward pass gradient
	# computation.
	state._post_backward_stream = state._device_handle.Stream()
	# Stream for pre-unshard logic, namely allocations and writes for CPU
	# offloading (H2D copy) and mixed precision (low precision cast).
	state._pre_unshard_stream = state._device_handle.Stream()
	# Default stream for computation
	state._default_stream = state._device_handle.current_stream()

[awgu]

We used to test that state._streams was shared. This adds back testing that each stream is shared.

[rohan-varma]

LGTM

[awgu]

@pytorchbot 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