diff --git a/deepspeed/runtime/zero/stage_1_and_2.py b/deepspeed/runtime/zero/stage_1_and_2.py index 922fe18b5184..95b2f9680ed5 100644 --- a/deepspeed/runtime/zero/stage_1_and_2.py +++ b/deepspeed/runtime/zero/stage_1_and_2.py @@ -139,6 +139,23 @@ def clear(self): self.copy_streams.clear() +class GradPartitionBufferGroup(list): + + def __init__(self, grad_buffers, flat_partition): + super().__init__(grad_buffers) + self.flat_partition = flat_partition + self._view_specs = [(buffer.storage_offset() - flat_partition.storage_offset(), buffer.shape, + getattr(buffer, "_zero_padding", False)) for buffer in grad_buffers] + + def move_flat_partition(self, device, non_blocking=False): + self.flat_partition.data = self.flat_partition.data.to(device, non_blocking=non_blocking) + for index, (offset, shape, is_padding) in enumerate(self._view_specs): + buffer = self.flat_partition.narrow(0, offset, math.prod(shape)).view(shape) + if is_padding: + buffer._zero_padding = True + self[index] = buffer + + class DeepSpeedZeroOptimizer(ZeROOptimizer): """ DeepSpeedZeroOptimizer designed to reduce the memory footprint @@ -1015,6 +1032,10 @@ def independent_gradient_partition_epilogue(self): if self.cpu_offload is False: for i, _ in enumerate(self.bit16_groups): + if self._use_compact_grad_partition_buffer(i): + if self.is_gradient_accumulation_boundary() and self.averaged_gradients.get(i) is None: + self.averaged_gradients[i] = self._build_compact_grad_partition_buffer(i) + continue if i not in self.all_grad_tensors or self.all_grad_tensors[i] is None: self.all_grad_tensors[i] = self.get_all_grad_tensors(self.params_in_partition[i], dtype=self.gradient_accumulation_dtype) @@ -1057,6 +1078,11 @@ def finalize_gradient_accumulation_boundary(self): self._finalize_cpu_offload_gradient_accumulation() return for i, _ in enumerate(self.bit16_groups): + if self._use_compact_grad_partition_buffer(i): + if self.averaged_gradients.get(i) is None: + self.averaged_gradients[i] = self._build_compact_grad_partition_buffer(i) + self.all_grad_tensors[i] = None + continue self.averaged_gradients[i] = self.get_flat_partition(self.params_in_partition[i], self.first_offset[i], self.partition_size[i], @@ -1068,12 +1094,15 @@ def finalize_gradient_accumulation_boundary(self): def _finalize_cpu_offload_gradient_accumulation(self): # Deferred boundary work for params reduced this window (matches managed offload; skips inactive params). - for group in self.params_in_partition: + for group_idx, group in enumerate(self.params_in_partition): for param in group: if not param.requires_grad: continue if self.get_param_id(param) not in self._offload_accumulated_param_ids: continue + if self._use_compact_grad_partition_buffer(group_idx): + self.copy_grads_in_partition(param, group_idx) + continue self._restore_cpu_offload_grad_to_gpu(param) self.set_norm_for_param_grad_in_gpu(param) self.update_offload_overflow_tracker_for_param_grad(param) @@ -1218,13 +1247,104 @@ def fill_grad_accum_attribute(self): def get_gradient_for_reduction(self, param): if self.use_grad_accum_attribute: - return param.grad_accum.to(self.dtype) if param.grad_accum is not None else None + return param.grad_accum if param.grad_accum is not None else None else: return param.grad + def _gradient_reduction_buffer_dtype(self): + return self.gradient_accumulation_dtype if self.use_grad_accum_attribute else self.dtype + def get_param_gradient_attribute(self, param): return param.grad_accum if self.use_grad_accum_attribute else param.grad + def _use_compact_grad_partition_buffer(self, group_idx): + gradient_accumulation_dtype = getattr(self, "gradient_accumulation_dtype", None) + if gradient_accumulation_dtype != torch.float32 or gradient_accumulation_dtype == self.dtype: + return False + params_in_partition = self.params_in_partition[group_idx] + if len(params_in_partition) > 0 and self._is_muon_group(params_in_partition): + return False + return True + + def _build_compact_grad_partition_buffer(self, group_idx): + partition_size = int(self.partition_size[group_idx]) + flat_partition = torch.zeros(partition_size, + dtype=self.gradient_accumulation_dtype, + device=get_accelerator().current_device_name()) + partition_id = dist.get_rank(group=self.real_dp_process_group[group_idx]) + + grad_buffers = [] + current_size = 0 + for tensor in self.params_in_partition[group_idx]: + param_id = self.get_param_id(tensor) + source_offset = int(self.grad_start_offset[group_idx][partition_id][param_id]) + dest_offset = int(self.grad_partition_insertion_offset[group_idx][partition_id][param_id]) + num_elements = int(min(tensor.numel() - source_offset, partition_size - dest_offset)) + + if dest_offset > current_size: + pad = flat_partition.narrow(0, current_size, dest_offset - current_size) + pad._zero_padding = True + grad_buffers.append(pad) + current_size = dest_offset + + if num_elements > 0: + grad_buffer = flat_partition.narrow(0, dest_offset, num_elements) + if source_offset == 0 and num_elements == tensor.numel(): + grad_buffer = grad_buffer.view(tensor.shape) + grad_buffers.append(grad_buffer) + current_size = dest_offset + num_elements + + if current_size < partition_size: + pad = flat_partition.narrow(0, current_size, partition_size - current_size) + pad._zero_padding = True + grad_buffers.append(pad) + + return GradPartitionBufferGroup(grad_buffers, flat_partition) + + def _get_compact_grad_partition_buffer(self, group_idx): + grad_group = self.averaged_gradients.get(group_idx) + if grad_group is not None and getattr(grad_group, "flat_partition", None) is not None: + return grad_group + + grad_group = self._build_compact_grad_partition_buffer(group_idx) + self.averaged_gradients[group_idx] = grad_group + return grad_group + + def _get_compact_grad_partition_slice(self, group_idx, partition_id, param_id, num_elements): + if not self._use_compact_grad_partition_buffer(group_idx): + return None + local_partition_id = dist.get_rank(group=self.real_dp_process_group[group_idx]) + if partition_id != local_partition_id: + return None + + dest_offset = int(self.grad_partition_insertion_offset[group_idx][partition_id][param_id]) + grad_group = self._get_compact_grad_partition_buffer(group_idx) + return grad_group.flat_partition.narrow(0, dest_offset, int(num_elements)) + + @staticmethod + def _add_to_grad_partition_buffer(output, source): + if output is not None: + output.add_(source.to(output.dtype)) + + def _add_full_reduction_to_compact_grad_partition(self, reduced_tensor, communication_data_type): + bucket = self.ipg_buckets[communication_data_type] + bucket_offset = 0 + for group_idx, param_idx_in_group, param_id in bucket.params: + param = self.bit16_groups[group_idx][param_idx_in_group] + if self._use_compact_grad_partition_buffer(group_idx): + partition_id = dist.get_rank(group=self.real_dp_process_group[group_idx]) + if param_id in self.grad_start_offset[group_idx][partition_id]: + source_offset = int(self.grad_start_offset[group_idx][partition_id][param_id]) + dest_offset = int(self.grad_partition_insertion_offset[group_idx][partition_id][param_id]) + num_elements = int(min(param.numel() - source_offset, + self.partition_size[group_idx] - dest_offset)) + if num_elements > 0: + source = reduced_tensor.narrow(0, bucket_offset + source_offset, num_elements) + output = self._get_compact_grad_partition_slice(group_idx, partition_id, param_id, + num_elements) + self._add_to_grad_partition_buffer(output, source) + bucket_offset += param.numel() + # Clear the tensor the reduction gradient attribute is pointing to def clear_grad_attribute(self, param): if self.use_grad_accum_attribute: @@ -1392,7 +1512,7 @@ def gradient_reduction_w_predivide(self, tensor, communication_data_type: torch. if communication_data_type != tensor.dtype and tensor is not tensor_to_allreduce: tensor.copy_(tensor_to_allreduce) - return tensor + return tensor_to_allreduce def allreduce_and_copy_with_multiple_ranks(self, small_bucket, @@ -1400,7 +1520,8 @@ def allreduce_and_copy_with_multiple_ranks(self, log=None, divide=True, process_group=None, - bucket_ranks=None): + bucket_ranks=None, + output_bucket=None): process_group = self.dp_process_group if process_group is None else process_group allreduced = self.allreduce_bucket(small_bucket, communication_data_type, @@ -1409,11 +1530,15 @@ def allreduce_and_copy_with_multiple_ranks(self, process_group=process_group) if self.overlap_comm and not get_accelerator().resolves_data_dependency(): allreduced.record_stream(self.reduction_stream) + synced_tensors = self.unflatten(allreduced, small_bucket) + if output_bucket is None: + output_bucket = [None] * len(small_bucket) local_rank = dist.get_rank(group=process_group) - for buf, synced, bucket_rank in zip(small_bucket, self.unflatten(allreduced, small_bucket), bucket_ranks): + for buf, synced, bucket_rank, output in zip(small_bucket, synced_tensors, bucket_ranks, output_bucket): copy_to_local_rank = local_rank in bucket_rank if isinstance(bucket_rank, frozenset) else local_rank == bucket_rank if copy_to_local_rank: + self._add_to_grad_partition_buffer(output, synced) buf.copy_(synced) if self.overlap_comm and not get_accelerator().resolves_data_dependency(): buf.record_stream(self.reduction_stream) @@ -1427,13 +1552,19 @@ def allreduce_and_scatter(self, process_group=None): small_bucket = [] small_bucket_ranks = [] + small_bucket_outputs = [] numel = 0 allreduce_sizes = [] for i, bucket_elem in enumerate(bucket): - rank, tensor = bucket_elem + if len(bucket_elem) == 2: + rank, tensor = bucket_elem + output = None + else: + rank, tensor, output = bucket_elem small_bucket.append(tensor) small_bucket_ranks.append(rank) + small_bucket_outputs.append(output) numel = numel + tensor.numel() if numel > numel_per_bucket: self.allreduce_and_copy_with_multiple_ranks(small_bucket, @@ -1441,9 +1572,11 @@ def allreduce_and_scatter(self, log=None, divide=divide, process_group=process_group, - bucket_ranks=small_bucket_ranks) + bucket_ranks=small_bucket_ranks, + output_bucket=small_bucket_outputs) small_bucket = [] small_bucket_ranks = [] + small_bucket_outputs = [] numel = 0 if len(small_bucket) > 0: @@ -1452,7 +1585,8 @@ def allreduce_and_scatter(self, log=None, divide=divide, process_group=process_group, - bucket_ranks=small_bucket_ranks) + bucket_ranks=small_bucket_ranks, + output_bucket=small_bucket_outputs) def average_tensor(self, tensor: torch.Tensor, communication_data_type: torch.dtype): if self._offload_gradient_safety_enabled: @@ -1481,7 +1615,8 @@ def average_tensor(self, tensor: torch.Tensor, communication_data_type: torch.dt with get_accelerator().stream(stream): self._record_gradient_stream(tensor, stream) if not self.reduce_scatter: - self.gradient_reduction_w_predivide(tensor, communication_data_type) + reduced_tensor = self.gradient_reduction_w_predivide(tensor, communication_data_type) + self._add_full_reduction_to_compact_grad_partition(reduced_tensor, communication_data_type) return # Accumulate destination ranks and bucket offsets for each gradient slice. @@ -1535,14 +1670,15 @@ def average_tensor(self, tensor: torch.Tensor, communication_data_type: torch.dt numel = partition_ids_w_offsets[idx + 1][1] - offset copy_ranks = muon_copy_ranks or frozenset((partition_id, )) + grad_partition_output = self._get_compact_grad_partition_slice(i, partition_id, param_id, numel) # Merge bucket ranges if they share the same reduction and consumers. - if (partition_id == prev_id and process_group == prev_process_group - and copy_ranks == prev_copy_ranks): - prev_pid, prev_size, prev_numel, _ = rank_and_offsets[-1] - rank_and_offsets[-1] = (prev_pid, prev_size, prev_numel + numel, copy_ranks) + if (grad_partition_output is None and partition_id == prev_id + and process_group == prev_process_group and copy_ranks == prev_copy_ranks): + prev_pid, prev_size, prev_numel, _, _ = rank_and_offsets[-1] + rank_and_offsets[-1] = (prev_pid, prev_size, prev_numel + numel, copy_ranks, None) else: - rank_and_offsets.append((partition_id, curr_size, numel, copy_ranks)) + rank_and_offsets.append((partition_id, curr_size, numel, copy_ranks, grad_partition_output)) real_dp_process_group.append(process_group) curr_size += numel prev_id, prev_process_group, prev_copy_ranks = partition_id, process_group, copy_ranks @@ -1551,7 +1687,7 @@ def average_tensor(self, tensor: torch.Tensor, communication_data_type: torch.dt tensor.div_(dist.get_world_size(group=self.dp_process_group) / float(self.sequence_parallel_size)) buckets = {} - for i, (dst, bucket_offset, numel, copy_ranks) in enumerate(rank_and_offsets): + for i, (dst, bucket_offset, numel, copy_ranks, grad_partition_output) in enumerate(rank_and_offsets): grad_slice = tensor.narrow(0, int(bucket_offset), int(numel)) process_group = real_dp_process_group[i] # Split Muon matrices require all-reduce even when multi-rank bucket all-reduce is disabled. @@ -1562,9 +1698,9 @@ def average_tensor(self, tensor: torch.Tensor, communication_data_type: torch.dt if bucket_key not in buckets: buckets[bucket_key] = [] if bucket_key[0] == "allreduce": - buckets[bucket_key].append((copy_ranks, grad_slice)) + buckets[bucket_key].append((copy_ranks, grad_slice, grad_partition_output)) else: - buckets[bucket_key].append(grad_slice) + buckets[bucket_key].append((grad_slice, grad_partition_output)) for bucket_key in buckets: if bucket_key[0] == "allreduce": @@ -1575,12 +1711,15 @@ def average_tensor(self, tensor: torch.Tensor, communication_data_type: torch.dt process_group=bucket_key[1]) else: _, dst, process_group = bucket_key - self.allreduce_no_retain(buckets[bucket_key], + grad_slices = [item[0] for item in buckets[bucket_key]] + grad_partition_outputs = [item[1] for item in buckets[bucket_key]] + self.allreduce_no_retain(grad_slices, communication_data_type, numel_per_bucket=self.reduce_bucket_size, rank=dst, divide=False, - process_group=process_group) + process_group=process_group, + output_bucket=grad_partition_outputs) ############################################################################## ############################# CPU Offload Methods############################# @@ -1983,7 +2122,39 @@ def _apply_muon_updates_cpu_offload(self): del full_grad, full_momentum, grad, param_momentum, update ############################################################################################ - def copy_grads_in_partition(self, param): + def _get_param_group_idx(self, param): + param_idx_in_group = getattr(param, "param_idx_in_group", None) + if param_idx_in_group is not None: + for group_idx, group in enumerate(self.bit16_groups): + if param_idx_in_group < len(group) and group[param_idx_in_group] is param: + return group_idx + + for group_idx, group in enumerate(self.bit16_groups): + for group_param in group: + if group_param is param: + return group_idx + + raise RuntimeError(f"Unable to find ZeRO parameter group for {debug_param2name(param)}") + + def copy_grads_in_partition(self, param, group_idx=None): + if group_idx is None: + group_idx = self._get_param_group_idx(param) + if self.cpu_offload and self._use_compact_grad_partition_buffer(group_idx): + param_id = self.get_param_id(param) + self._offload_accumulated_param_ids.add(param_id) + if self.is_gradient_accumulation_boundary(): + [_, _, dest_offset, num_elements] = self.grad_position[param_id] + partition_id = dist.get_rank(group=self.real_dp_process_group[group_idx]) + source = self._get_compact_grad_partition_slice(group_idx, partition_id, param_id, num_elements) + self.norm_for_param_grads[param_id] = source.data.to(get_norm_dtype()).norm(2) + self.update_offload_overflow_tracker(source) + destination = self.single_partition_of_fp32_groups[group_idx].grad.view(-1).narrow( + 0, dest_offset, num_elements) + destination.copy_(source.to(destination.dtype), non_blocking=True) + self._record_gradient_stream(source, get_accelerator().current_stream()) + self._record_offload_copy() + self.clear_grad_attribute(param) + return if self.cpu_offload: # Accumulate when there were prior backwards in this step (restore from # CPU buffer) or more will follow (save to CPU buffer). Skipping only @@ -2002,6 +2173,10 @@ def copy_grads_in_partition(self, param): self.async_inplace_copy_grad_to_fp32_buffer_from_gpu(param) return + if self._use_compact_grad_partition_buffer(group_idx): + if self.use_grad_accum_attribute: + self.clear_grad_attribute(param) + return #print(f"ID {self.get_param_id(param)} grad norm {param.grad.norm()}") if self.grads_in_partition is None: self.grads_in_partition_offset = 0 @@ -2050,7 +2225,13 @@ def reduce_ipg_grads(self, comm_dtype=None): else: self.average_tensor(bucket.buffer[bucket.index].narrow(0, 0, bucket.elements), comm_dtype) else: - self.buffered_reduce_fallback(None, bucket.grads, comm_dtype, elements_per_buffer=bucket.elements) + if any(self._use_compact_grad_partition_buffer(group_idx) for group_idx, _, _ in bucket.params): + self._buffered_reduce_fallback_preserving_fp32(bucket.grads, + bucket.params, + comm_dtype, + elements_per_buffer=bucket.elements) + else: + self.buffered_reduce_fallback(None, bucket.grads, comm_dtype, elements_per_buffer=bucket.elements) if self.overlap_comm: stream = self.reduction_stream @@ -2084,10 +2265,10 @@ def reduce_ipg_grads(self, comm_dtype=None): else: self.clear_grad_attribute(param) elif self.contiguous_gradients: - self.copy_grads_in_partition(param) + self.copy_grads_in_partition(param, group_idx) else: # zero stage 1 - partition only optimizer state if self.contiguous_gradients and self.is_param_in_current_partition[param_id]: - self.copy_grads_in_partition(param) + self.copy_grads_in_partition(param, group_idx) # Empty and oversized reductions do not consume the contiguous buffer. if (self._offload_gradient_safety_enabled and 0 < bucket.elements <= self.reduce_bucket_size and not get_accelerator().resolves_data_dependency()): @@ -2206,7 +2387,7 @@ def allreduce_bucket(self, if rank is None or rank == dist.get_rank(group=process_group): tensor.copy_(tensor_to_allreduce) - return tensor + return tensor_to_allreduce def _clear_previous_reduced_grads(self): for dtype in self.previous_reduced_grads: @@ -2221,7 +2402,8 @@ def allreduce_and_copy(self, rank=None, log=None, divide=True, - process_group=None): + process_group=None, + output_bucket=None): process_group = self.dp_process_group if process_group is None else process_group if self.overlap_comm: if not get_accelerator().resolves_data_dependency(): @@ -2243,8 +2425,12 @@ def allreduce_and_copy(self, ) if self.overlap_comm and not get_accelerator().resolves_data_dependency(): allreduced.record_stream(stream) - if rank is None or rank == dist.get_rank(group=self.dp_process_group): - for buf, synced in zip(small_bucket, self.unflatten(allreduced, small_bucket)): + if rank is None or rank == dist.get_rank(group=process_group): + synced_tensors = self.unflatten(allreduced, small_bucket) + if output_bucket is None: + output_bucket = [None] * len(small_bucket) + for buf, synced, output in zip(small_bucket, synced_tensors, output_bucket): + self._add_to_grad_partition_buffer(output, synced) buf.copy_(synced) if self.overlap_comm and not get_accelerator().resolves_data_dependency(): buf.record_stream(stream) @@ -2258,11 +2444,16 @@ def allreduce_no_retain( log=None, divide=True, process_group=None, + output_bucket=None, ): small_bucket = [] + small_output_bucket = [] numel = 0 - for tensor in bucket: + if output_bucket is None: + output_bucket = [None] * len(bucket) + for tensor, output in zip(bucket, output_bucket): small_bucket.append(tensor) + small_output_bucket.append(output) numel = numel + tensor.numel() if numel > numel_per_bucket: self.allreduce_and_copy(small_bucket, @@ -2270,8 +2461,10 @@ def allreduce_no_retain( rank=rank, log=None, divide=divide, - process_group=process_group) + process_group=process_group, + output_bucket=small_output_bucket) small_bucket = [] + small_output_bucket = [] numel = 0 if len(small_bucket) > 0: @@ -2280,10 +2473,65 @@ def allreduce_no_retain( rank=rank, log=log, divide=divide, - process_group=process_group) + process_group=process_group, + output_bucket=small_output_bucket) # allows using reduction of gradients instead of using all_reduce + def _allreduce_noncontiguous_bucket_preserving_fp32(self, grads, grad_params, communication_data_type): + if self.overlap_comm: + if not get_accelerator().resolves_data_dependency(): + get_accelerator().synchronize() + self._clear_previous_reduced_grads() + stream = self.reduction_stream + else: + stream = get_accelerator().current_stream() + + with get_accelerator().stream(stream): + allreduced = self.allreduce_bucket(grads, communication_data_type) + if self.overlap_comm and not get_accelerator().resolves_data_dependency(): + allreduced.record_stream(stream) + reduced_grads = self.unflatten(allreduced, grads) + for grad, reduced_grad, (group_idx, _, param_id) in zip(grads, reduced_grads, grad_params): + partition_id = dist.get_rank(group=self.real_dp_process_group[group_idx]) + if param_id in self.grad_start_offset[group_idx][partition_id]: + param = self.param_dict[param_id] + source_offset = int(self.grad_start_offset[group_idx][partition_id][param_id]) + dest_offset = int(self.grad_partition_insertion_offset[group_idx][partition_id][param_id]) + num_elements = int(min(param.numel() - source_offset, + self.partition_size[group_idx] - dest_offset)) + output = self._get_compact_grad_partition_slice(group_idx, partition_id, param_id, num_elements) + source = reduced_grad.view(-1).narrow(0, source_offset, num_elements) + self._add_to_grad_partition_buffer(output, source) + grad.copy_(reduced_grad) + if self.overlap_comm and not get_accelerator().resolves_data_dependency(): + grad.record_stream(stream) + + def _buffered_reduce_fallback_preserving_fp32(self, grads, grad_params, communication_data_type, + elements_per_buffer): + buckets = {} + for grad, grad_param in zip(grads, grad_params): + dtype_bucket = buckets.setdefault(grad.dtype, []) + dtype_bucket.append((grad, grad_param)) + + for dtype_bucket in buckets.values(): + bucket_grads = [] + bucket_params = [] + numel = 0 + for grad, grad_param in dtype_bucket: + bucket_grads.append(grad) + bucket_params.append(grad_param) + numel += grad.numel() + if numel > elements_per_buffer: + self._allreduce_noncontiguous_bucket_preserving_fp32(bucket_grads, bucket_params, + communication_data_type) + bucket_grads = [] + bucket_params = [] + numel = 0 + if bucket_grads: + self._allreduce_noncontiguous_bucket_preserving_fp32(bucket_grads, bucket_params, + communication_data_type) + def buffered_reduce_fallback(self, rank, grads, @@ -2842,6 +3090,8 @@ def step(self, closure=None): bit16_partition_buffer = self.param_buffer_of_bit16_for_cpu_offload_groups[i] bit16_partition_buffer.data.copy_(fp32_partition.data) bit16_partitions[partition_id].data.copy_(bit16_partition_buffer.data, non_blocking=True) + self._release_preflattened_grad_buffers(i) + self.averaged_gradients[i] = None self.timers(OPTIMIZER_STEP_TIMER).stop() else: @@ -3007,9 +3257,10 @@ def setup_buckets(self): for _, bucket in self.ipg_buckets.items(): bucket.buffer.clear() - # Buffer's dtype is the same as the dtype of optimizer, not dtype for autocast + # ZeRO-1 stages the fp32 accumulator; ZeRO-2 keeps low-precision staging + # and stores the fp32 communication result in the owned partition. buf_0 = torch.empty(int(self.reduce_bucket_size), - dtype=self.dtype, + dtype=self._gradient_reduction_buffer_dtype(), device=get_accelerator().current_device_name()) bucket.buffer.append(buf_0) bucket.index = 0 @@ -3018,7 +3269,7 @@ def setup_buckets(self): if self.overlap_comm: for _, bucket in self.ipg_buckets.items(): buf_1 = torch.empty(int(self.reduce_bucket_size), - dtype=self.dtype, + dtype=self._gradient_reduction_buffer_dtype(), device=get_accelerator().current_device_name()) bucket.buffer.append(buf_1) @@ -3653,12 +3904,13 @@ def needs_offload(target): for group_idx in self.averaged_gradients: grad_list = self.averaged_gradients.get(group_idx) if grad_list is not None: - for grad_tensor in grad_list: - if grad_tensor is not None and grad_tensor.device.type != device: - # Key insight: We only move the underlying data storage (`.data`) to the target device. - # The Python tensor object and the dictionary structure (`self.averaged_gradients`) - # remain intact, preserving the references needed for reloading. - grad_tensor.data = grad_tensor.data.to(device, non_blocking=non_blocking) + move_flat_partition = getattr(grad_list, "move_flat_partition", None) + if callable(move_flat_partition): + move_flat_partition(device, non_blocking=non_blocking) + else: + for grad_tensor in grad_list: + if grad_tensor is not None and grad_tensor.device.type != device: + grad_tensor.data = grad_tensor.data.to(device, non_blocking=non_blocking) self.offloaded_states.add(OffloadStateTypeEnum.lp_grads) @@ -3722,14 +3974,16 @@ def reload_states(self, non_blocking: bool = False): # Reload Partitioned Gradients (LP Grads) if OffloadStateTypeEnum.lp_grads in self.offloaded_states: - # Since we preserved the `self.averaged_gradients` structure during offload, - # we can now iterate through it again. The tensors within currently point to CPU data. for group_idx in self.averaged_gradients: grad_list = self.averaged_gradients.get(group_idx) if grad_list is not None: - for grad_tensor in grad_list: - if grad_tensor is not None and grad_tensor.device.type != device: - grad_tensor.data = grad_tensor.data.to(device, non_blocking=non_blocking) + move_flat_partition = getattr(grad_list, "move_flat_partition", None) + if callable(move_flat_partition): + move_flat_partition(device, non_blocking=non_blocking) + else: + for grad_tensor in grad_list: + if grad_tensor is not None and grad_tensor.device.type != device: + grad_tensor.data = grad_tensor.data.to(device, non_blocking=non_blocking) self.offloaded_states.remove(OffloadStateTypeEnum.lp_grads) diff --git a/deepspeed/utils/tensor_fragment.py b/deepspeed/utils/tensor_fragment.py index f7c7c0995594..296517ea44a2 100644 --- a/deepspeed/utils/tensor_fragment.py +++ b/deepspeed/utils/tensor_fragment.py @@ -217,6 +217,10 @@ def safe_get_full_grad(param): Returns: Union[torch.Tensor, None]: A tensor on accelerator device """ + # BF16_Optimizer keeps its accumulation buffer separately while retaining the low-precision staging gradient. + if getattr(param, '_hp_grad', None) is not None: + return param.get_full_hp_grad() + if param.grad is not None: return param.grad diff --git a/tests/unit/runtime/zero/test_zero_tensor_fragment.py b/tests/unit/runtime/zero/test_zero_tensor_fragment.py index 1214a8247acb..83c42cdad3bf 100644 --- a/tests/unit/runtime/zero/test_zero_tensor_fragment.py +++ b/tests/unit/runtime/zero/test_zero_tensor_fragment.py @@ -19,7 +19,7 @@ from deepspeed.utils import safe_set_local_fp32_param, safe_set_local_grad, safe_set_local_optimizer_state from deepspeed.utils import safe_update_full_grad_vectorized from deepspeed.utils.mixed_precision_linkage import link_hp_params -from deepspeed.runtime.zero.offload_config import OffloadDeviceEnum +from deepspeed.runtime.zero.offload_config import OffloadDeviceEnum, OffloadStateTypeEnum from deepspeed.ops.aio import AsyncIOBuilder from deepspeed.accelerator import get_accelerator from deepspeed.runtime.swap_tensor import MIN_SWAPPABLE_BYTES @@ -99,6 +99,16 @@ def forward(self, x, y): return self.cel(x, y) +class ScaleModel(torch.nn.Module): + + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.ones(4)) + + def forward(self, x): + return (self.weight * x).sum() + + def run_fragmented_model(model, config_dict, hidden_dim, dtype, validate_after_bwd, validate_after_step): model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict) data_loader = random_dataloader(model=model, @@ -229,6 +239,109 @@ def test_bf16_optimizer_fragments(self, frozen_weights): run_fragmented_model(model, config_dict, hidden_dim, torch.bfloat16, validate_after_bwd, validate_after_step) +class TestZeroBf16Fp32GradAccum(DistributedTest): + world_size = 2 + + @pytest.mark.parametrize( + 'zero_stage, zero_options', + [ + pytest.param(1, {}, id="stage1"), + pytest.param(2, {}, id="stage2"), + pytest.param(2, {"offload_optimizer": { + "device": "cpu" + }}, id="stage2-cpu-offload"), + pytest.param(2, {"contiguous_gradients": False}, id="stage2-noncontiguous"), + pytest.param(3, {}, id="stage3"), + ], + ) + def test_safe_get_full_grad_preserves_fp32_accumulated_average(self, zero_stage, zero_options): + # Catches storing or accumulating fp32 communication results in a bf16 partition. + if not get_accelerator().is_bf16_supported(): + pytest.skip("requires bf16") + + config_dict = { + "train_micro_batch_size_per_gpu": 1, + "gradient_accumulation_steps": 4, + "steps_per_print": 1, + "communication_data_type": "fp32", + "optimizer": { + "type": "Adam", + "params": { + "lr": 1e-6, + "torch_adam": True + } + }, + "bf16": { + "enabled": True + }, + "data_types": { + "grad_accum_dtype": "fp32" + }, + "zero_optimization": { + "stage": zero_stage, + **zero_options, + } + } + + model = ScaleModel() + engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict) + + rank = dist.get_rank() + local_grad = 1.0 if rank == 0 else 1.0078125 + x = torch.full((4, ), local_grad, dtype=torch.bfloat16, device=engine.device) + for micro_step in range(engine.gradient_accumulation_steps()): + engine.backward(engine(x)) + if not engine.is_gradient_accumulation_boundary(): + engine.step() + + grad = safe_get_full_grad(engine.module.weight) + expected = torch.full((4, ), 1.00390625, dtype=torch.float32, device=engine.device) + assert grad.dtype == torch.float32 + torch.testing.assert_close(grad, expected, rtol=0, atol=0) + + engine.destroy() + + def test_dynamic_gradient_offload_preserves_optimizer_input(self): + if not get_accelerator().is_bf16_supported(): + pytest.skip("requires bf16") + + config_dict = { + "train_micro_batch_size_per_gpu": 1, + "communication_data_type": "fp32", + "zero_allow_untested_optimizer": True, + "optimizer": { + "type": "SGD", + "params": { + "lr": 1.0 + } + }, + "bf16": { + "enabled": True + }, + "data_types": { + "grad_accum_dtype": "fp32" + }, + "zero_optimization": { + "stage": 2, + }, + } + + model = ScaleModel() + engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict) + x = torch.arange(1, 5, dtype=torch.bfloat16, device=engine.device) + engine.backward(engine(x)) + engine.offload_states(include=[OffloadStateTypeEnum.lp_grads], pin_memory=False) + engine.reload_states() + + weight_before = engine.module.weight.detach().clone() + replacement_grad = torch.zeros(4, dtype=torch.float32, device=engine.device) + safe_set_full_grad(engine.module.weight, replacement_grad) + engine.step() + + torch.testing.assert_close(engine.module.weight, weight_before, rtol=0, atol=0) + engine.destroy() + + def create_random_values(model, key_list, group, grad_dtype): param_values = {} for n, lp in model.named_parameters(): diff --git a/tests/unit/v1/zero/test_zero2_gradient_safety.py b/tests/unit/v1/zero/test_zero2_gradient_safety.py index c832949cabc4..a6a83101e0da 100644 --- a/tests/unit/v1/zero/test_zero2_gradient_safety.py +++ b/tests/unit/v1/zero/test_zero2_gradient_safety.py @@ -87,6 +87,34 @@ def backward_values(opt, param, values, boundaries): opt._wait_for_offload_copies() +def test_allreduce_and_scatter_accepts_legacy_rank_tensor_buckets(): + opt = zero.DeepSpeedZeroOptimizer.__new__(zero.DeepSpeedZeroOptimizer) + captured = {} + + def capture(small_bucket, + communication_data_type, + log=None, + divide=True, + process_group=None, + bucket_ranks=None, + output_bucket=None): + captured["small_bucket"] = small_bucket + captured["bucket_ranks"] = bucket_ranks + captured["output_bucket"] = output_bucket + + opt.allreduce_and_copy_with_multiple_ranks = capture + tensor = torch.ones(2) + output = torch.empty_like(tensor) + + opt.allreduce_and_scatter([(0, tensor), (1, tensor, output)], torch.float32) + + assert captured["small_bucket"][0] is tensor + assert captured["small_bucket"][1] is tensor + assert captured["bucket_ranks"] == [0, 1] + assert captured["output_bucket"][0] is None + assert captured["output_bucket"][1] is output + + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) @pytest.mark.parametrize("cpu_offload", [False, True]) @pytest.mark.parametrize("bucket_size", [4, 8, 16]) @@ -224,7 +252,7 @@ def test_only_buffer_consumers_replace_reuse_event(monkeypatch, bucket_size, ele opt.reduce_bucket_size = bucket_size opt.is_param_in_current_partition = {0: True} opt.average_tensor = lambda *args: None - opt.copy_grads_in_partition = lambda param: None + opt.copy_grads_in_partition = lambda param, group_idx=None: None bucket = opt.ipg_buckets[torch.float32] bucket.elements = elements previous, consumer = FakeEvent(), FakeStream()