Skip to content

Commit

Permalink
fix ds-inference without policy (#2247)
Browse files Browse the repository at this point in the history
Co-authored-by: Jeff Rasley <jerasley@microsoft.com>
  • Loading branch information
RezaYazdaniAminabadi and jeffra committed Aug 23, 2022
1 parent fae896e commit c35bfe8
Showing 1 changed file with 4 additions and 3 deletions.
7 changes: 4 additions & 3 deletions deepspeed/module_inject/replace_module.py
Original file line number Diff line number Diff line change
Expand Up @@ -653,11 +653,10 @@ def _replace(child, name, conv_linear_layer):
new_bias = torch.empty((weight_shape[0]),
device=child.weight.device,
dtype=child.weight.dtype)

if z_inference:
with deepspeed.zero.GatheredParameters(child.bias, modifier_rank=0):
new_bias.data.copy_(child.bias.data)
else:
elif child.bias:
new_bias.data.copy_(child.bias.data)
return LinearAllreduce(data, child.bias if child.bias is None else \
torch.nn.parameter.Parameter(new_bias.to(torch.cuda.current_device())), mp_group)
Expand Down Expand Up @@ -704,7 +703,9 @@ def _slice_embedding(child, name, conv_linear_layer):
child.weight.shape[1] // mp_size),
device=child.weight.device,
dtype=child.weight.dtype)
data = mp_replace.copy(new_weight, child.weight.ds_tensor.data)
data = mp_replace.copy(new_weight,
child.weight.ds_tensor.data if hasattr(child.weight, 'ds_tensor') else \
child.weight.data)
new_embedding = nn.Embedding(child.weight.shape[0],
child.weight.shape[1] // mp_size)
new_embedding.weight.data.copy_(data)
Expand Down

0 comments on commit c35bfe8

Please sign in to comment.