Feature/group offload pinning #12747
Conversation
Aki-07
force-pushed
the
feature/group-offload-pinning
branch
from
7e50d90 to
3b3813d
Compare
|
Thanks for your PR. However, it's being worked on in #12721. |
|
Could we resolve conflicts so that it's a bit easier to review? Seems like there's some overlap from #12692. |
Aki-07
force-pushed
the
feature/group-offload-pinning
branch
from
6d96002 to
33d8b52
Compare
|
Done! Rebased on latest main and resolved conflicts with #12692. Should be much cleaner to review now. |
| should_synchronize = ( | ||
| not self.group.onload_self and self.group.stream is not None and not should_onload_next_group | ||
| ) |
There was a problem hiding this comment.
Even with non_blocking=True, if a previous group onloaded this one on a side stream, we need a sync before the default stream uses the weights or we risk reading half-copied tensors. I’ve limited the sync to the record_stream=False case, when record_stream=True the tensors are tied to the consumer stream so we can safely skip the sync.
| if len(tensors) == 0: | ||
| return True |
There was a problem hiding this comment.
This means the group is empty. Why would we return True for this?
There was a problem hiding this comment.
Agreed, that was misleading. Now an empty group returns False so a ‘pinned’ empty group will still onload instead of claiming it’s already on device.
| normalized_pin_groups = pin_groups | ||
| if isinstance(pin_groups, str): | ||
| normalized_pin_groups = pin_groups.lower() | ||
| if normalized_pin_groups not in {"first_last", "all"}: | ||
| raise ValueError("`pin_groups` must be one of `None`, 'first_last', 'all', or a callable.") | ||
| elif pin_groups is not None and not callable(pin_groups): | ||
| raise ValueError("`pin_groups` must be one of `None`, 'first_last', 'all', or a callable.") | ||
|
|
||
| pin_groups = normalized_pin_groups |
There was a problem hiding this comment.
(nit): would prefer to have a small utility function: _normalize_pin_groups().
There was a problem hiding this comment.
Added a _normalize_pin_groups helper fn()
| for name, submodule in module.named_children(): | ||
| if not isinstance(submodule, (torch.nn.ModuleList, torch.nn.Sequential)): | ||
| # Check if this is an explicitly defined block module | ||
| if name in block_modules: |
There was a problem hiding this comment.
| if name in block_modules: | |
| if block_modules and name in block_modules: |
| block_modules: Optional[List[str]] = None | ||
| exclude_kwargs: Optional[List[str]] = None | ||
| module_prefix: Optional[str] = "" | ||
| pin_groups: Optional[Union[str, Callable]] = None |
There was a problem hiding this comment.
This seems like a breaking change. Could you please elaborate?
There was a problem hiding this comment.
Thanks for flagging, have fixed to default it. Sorry for the oversight
| tensor.data = source_tensor.to(self.onload_device, non_blocking=self.non_blocking) | ||
| if self.record_stream: | ||
| tensor.data.record_stream(default_stream) | ||
| tensor.data.record_stream(self._torch_accelerator_module.current_stream()) |
There was a problem hiding this comment.
Could you elaborate on this change?
#12721 explains why it's the way it is.
There was a problem hiding this comment.
Fixed to the correct behavior from #12721, record the tensor on the consumer/default stream ( captured before entering the transfer stream ) so its lifetime is tied to the forward stream
|
Thank you for the initial comment! We are working on the solutions right now |
3 participants
What does this PR do?
Fixes #11966
Before submitting
documentation guidelines, and
here are tips on formatting docstrings.
Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.
@sayakpaul