Skip to content

Distributed optimizer support for contiguous param buffer with FP8 params #1749

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
crcrpar merged 3 commits into from
Nov 20, 2023
+34 −17
Merged

Distributed optimizer support for contiguous param buffer with FP8 params #1749

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
40500a3
Debug distopt contiguous param buffers with uint8 param all-gathers
timmoon10 Oct 31, 2023
1696bd4
Add test
timmoon10 Nov 17, 2023
42d5d8d
Avoid temporary buffer for param shard in optim step if possible
timmoon10 Nov 19, 2023
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
45 changes: 28 additions & 17 deletions apex/contrib/optimizers/distributed_fused_adam.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1007,7 +1007,7 @@ def init_param_buffer(self) -> None:
if param_buffer_view.dtype != param.dtype:
raise RuntimeError(
f"Attempted to change a parameter with dtype={param.dtype} "
f"into a buffer view with dtype={param_view_buffer.dtype}"
f"into a buffer view with dtype={param_buffer_view.dtype}"
)
param_flat_views.append(param.detach().view(-1))
param_buffer_views.append(param_buffer_view)
Expand Down Expand Up @@ -2210,8 +2210,23 @@ def step(
params_bucket = self.ParameterBucket()
state_bucket = self.state["buckets"][bucket_id]
shard_size = state_bucket.shard_size
dtype = state_bucket.dtype
param_sync_dtype = state_bucket.param_sync_dtype
if self.contiguous_param_buffer:
if not param_sync_dtype.is_floating_point:
# Make sure param shard buffer is floating-point
overlap_first_bucket = False
if (
state_bucket.params_shard is not None
and dtype.is_floating_point
):
params_bucket.params_shard = state_bucket.params_shard
else:
params_bucket.params_shard = torch.empty(
[shard_size],
dtype=self.dtype,
device=self.device,
)
elif self.contiguous_param_buffer:
# Construct view into contiguous param buffer
if not self._param_buffers:
self.init_param_buffer()
Expand All @@ -2225,23 +2240,19 @@ def step(
params_bucket.params_shard = params_bucket.params_bucket[
bucket_start:bucket_end
]
elif not param_sync_dtype.is_floating_point:
# Allocate temporary buffer for param shard
# Note: Adam kernel only supports floating-point
# dtypes.
params_bucket.params_shard = torch.empty(
[shard_size],
dtype=self.dtype,
device=self.device,
)
overlap_first_bucket = False
else:
# Allocate param shard buffer
params_bucket.params_shard = torch.empty(
[shard_size],
dtype=param_sync_dtype,
device=self.device,
)
if (
state_bucket.params_shard is not None
and dtype == param_sync_dtype
):
params_bucket.params_shard = state_bucket.params_shard
else:
params_bucket.params_shard = torch.empty(
[shard_size],
dtype=param_sync_dtype,
device=self.device,
)
self._params_buckets[bucket_id] = params_bucket

# Apply optimizer step
Expand Down
6 changes: 6 additions & 0 deletions apex/contrib/test/optimizers/test_dist_adam.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -296,6 +296,12 @@ def test_matches_pytorch_int64_param_sync(self):
param_sync_dtype=torch.int64,
)

def test_matches_pytorch_int64_param_sync_contiguous_buffers(self):
self.test_matches_pytorch(
param_sync_dtype=torch.int64,
contiguous_buffers=True,
)

def test_matches_pytorch_uint8_param_sync(self):
self.test_matches_pytorch(
rtol=0.5,
Expand Down