Skip to content

GitLab

"...gmock/git@developer.sourcefind.cn:yangql/googletest.git" did not exist on "c9ccac7cb7345901884aabf5d1a786cfa6e2f397"
Unverified Commit 11c5d588 authored by Kirthi Shankar Sivamani's avatar Kirthi Shankar Sivamani Committed by GitHub
Browse files

Remove buffer registration for FSDP like cases (#318) 



Remove extra buffers
Signed-off-by: default avatarKirthi Shankar Sivamani <ksivamani@nvidia.com>
parent 8c3110d1
Hide whitespace changes
Inline Side-by-side
Showing with 28 additions and 50 deletions
transformer_engine/pytorch/module/layernorm_linear.py
View file @ 11c5d588
...@@ -696,13 +696,10 @@ class LayerNormLinear(TransformerEngineBaseModule): ...@@ -696,13 +696,10 @@ class LayerNormLinear(TransformerEngineBaseModule):
self.reset_layer_norm_parameters() self.reset_layer_norm_parameters()
if not skip_weight_param_allocation: if not skip_weight_param_allocation:
self.register_buffer("weight_tensor", self.weight_tensor = torch.empty(
torch.empty( self.out_features, self.in_features,
self.out_features, device=torch.cuda.current_device(),
self.in_features, dtype=params_dtype)
device=torch.cuda.current_device(),
dtype=params_dtype),
persistent=False)
initialize_affine_weight_gpu( initialize_affine_weight_gpu(
self.weight_tensor, self.weight_tensor,
...@@ -713,17 +710,13 @@ class LayerNormLinear(TransformerEngineBaseModule): ...@@ -713,17 +710,13 @@ class LayerNormLinear(TransformerEngineBaseModule):
) )
if self.use_bias: if self.use_bias:
self.register_buffer("bias_tensor", self.bias_tensor = torch.empty(
torch.empty( self.out_features,
self.out_features, device=torch.cuda.current_device(),
device=torch.cuda.current_device(), dtype=params_dtype)
dtype=params_dtype),
persistent=False)
else: else:
self.register_buffer("bias_tensor", self.bias_tensor = torch.Tensor().to(dtype=params_dtype,
torch.Tensor().to(dtype=params_dtype, device=torch.cuda.current_device())
device=torch.cuda.current_device()),
persistent=False)
with torch.no_grad(): with torch.no_grad():
self.bias_tensor.zero_() self.bias_tensor.zero_()
...@@ -760,10 +753,8 @@ class LayerNormLinear(TransformerEngineBaseModule): ...@@ -760,10 +753,8 @@ class LayerNormLinear(TransformerEngineBaseModule):
bname, Parameter(self.bias_tensor[i * split_size : (i+1) * split_size]) bname, Parameter(self.bias_tensor[i * split_size : (i+1) * split_size])
) )
else: else:
self.register_buffer(bname, setattr(self, bname, torch.Tensor().to(dtype=params_dtype,
torch.Tensor().to(dtype=params_dtype, device=torch.cuda.current_device()))
device=torch.cuda.current_device()),
persistent=False)
if parallel_mode == "column": if parallel_mode == "column":
set_tensor_model_parallel_attributes(getattr(self, bname), True, 0, 1) set_tensor_model_parallel_attributes(getattr(self, bname), True, 0, 1)
......
transformer_engine/pytorch/module/layernorm_mlp.py
View file @ 11c5d588
...@@ -1048,10 +1048,8 @@ class LayerNormMLP(TransformerEngineBaseModule): ...@@ -1048,10 +1048,8 @@ class LayerNormMLP(TransformerEngineBaseModule):
) )
set_tensor_model_parallel_attributes(self.fc1_bias, True, 0, 1) set_tensor_model_parallel_attributes(self.fc1_bias, True, 0, 1)
else: else:
self.register_buffer("fc1_bias", self.fc1_bias = torch.Tensor().to(dtype=params_dtype,
torch.Tensor().to(dtype=params_dtype, device=torch.cuda.current_device())
device=torch.cuda.current_device()),
persistent=False)
with torch.no_grad(): with torch.no_grad():
self.fc1_bias.zero_() self.fc1_bias.zero_()
...@@ -1082,10 +1080,8 @@ class LayerNormMLP(TransformerEngineBaseModule): ...@@ -1082,10 +1080,8 @@ class LayerNormMLP(TransformerEngineBaseModule):
) )
) )
else: else:
self.register_buffer("fc2_bias", self.fc2_bias = torch.Tensor().to(dtype=params_dtype,
torch.Tensor().to(dtype=params_dtype, device=torch.cuda.current_device())
device=torch.cuda.current_device()),
persistent=False)
# For RPL, bias has to be added after TP collectives # For RPL, bias has to be added after TP collectives
# So it cannot be fused with the GEMM # So it cannot be fused with the GEMM
......
transformer_engine/pytorch/module/linear.py
View file @ 11c5d588
...@@ -559,13 +559,10 @@ class Linear(TransformerEngineBaseModule): ...@@ -559,13 +559,10 @@ class Linear(TransformerEngineBaseModule):
self.sequence_parallel = (self.tp_size > 1) and sequence_parallel self.sequence_parallel = (self.tp_size > 1) and sequence_parallel
if not skip_weight_param_allocation: if not skip_weight_param_allocation:
self.register_buffer("weight_tensor", self.weight_tensor = torch.empty(
torch.empty( self.out_features, self.in_features,
self.out_features, device=torch.cuda.current_device(),
self.in_features, dtype=params_dtype)
device=torch.cuda.current_device(),
dtype=params_dtype),
persistent=False)
initialize_affine_weight_gpu( initialize_affine_weight_gpu(
self.weight_tensor, self.weight_tensor,
...@@ -576,17 +573,13 @@ class Linear(TransformerEngineBaseModule): ...@@ -576,17 +573,13 @@ class Linear(TransformerEngineBaseModule):
) )
if self.use_bias: if self.use_bias:
self.register_buffer("bias_tensor", self.bias_tensor = torch.empty(
torch.empty( self.out_features,
self.out_features, device=torch.cuda.current_device(),
device=torch.cuda.current_device(), dtype=params_dtype)
dtype=params_dtype),
persistent=False)
else: else:
self.register_buffer("bias_tensor", self.bias_tensor = torch.Tensor().to(dtype=params_dtype,
torch.Tensor().to(dtype=params_dtype, device=torch.cuda.current_device())
device=torch.cuda.current_device()),
persistent=False)
with torch.no_grad(): with torch.no_grad():
self.bias_tensor.zero_() self.bias_tensor.zero_()
...@@ -623,10 +616,8 @@ class Linear(TransformerEngineBaseModule): ...@@ -623,10 +616,8 @@ class Linear(TransformerEngineBaseModule):
bname, Parameter(self.bias_tensor[i * split_size : (i+1) * split_size]) bname, Parameter(self.bias_tensor[i * split_size : (i+1) * split_size])
) )
else: else:
self.register_buffer(bname, setattr(self, bname, torch.Tensor().to(dtype=params_dtype,
torch.Tensor().to(dtype=params_dtype, device=torch.cuda.current_device()))
device=torch.cuda.current_device()),
persistent=False)
if parallel_mode == "column": if parallel_mode == "column":
set_tensor_model_parallel_attributes(getattr(self, bname), True, 0, 1) set_tensor_model_parallel_attributes(getattr(self, bname), True, 0, 1)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment