QKV parameters unfused path fixes and optimization

transformer_engine/pytorch/module.py

@@ -7,7 +7,7 @@
 import pickle
 import warnings
 from abc import ABC, abstractmethod
-from typing import Union, Optional, Callable, Tuple, Dict, Any, Mapping
+from typing import Union, Optional, Callable, Tuple, Dict, Any, Mapping, List
 from functools import partial
 from contextlib import contextmanager
 
@@ -135,6 +135,42 @@ def _prepare_backward(fp8: bool,
  delete_key_from_amax_buffer(forward=False)
 
 
+class _NoopCat(torch.autograd.Function):
+ """This class is a no-op replacement for `torch.cat`."""
+
+ @staticmethod
+ def forward(ctx,
+ full_param_buffer: torch.Tensor,
+ *params_split: Tuple[torch.Tensor, ...],
+ ) -> torch.Tensor:
+ assert not full_param_buffer.requires_grad, "Buffers should not require gradient"
+ assert (
+ full_param_buffer.shape[0] % len(params_split) == 0
+ ), "Dimensions not compatible for concatenation"
+
+ param_temp = full_param_buffer.new()
+ param_temp.set_(full_param_buffer.storage(),
+ full_param_buffer.storage_offset(),
+ full_param_buffer.size(),
+ full_param_buffer.stride())
+ param_temp.requires_grad = True
+
+ ctx.save_for_backward(full_param_buffer, *params_split)
+ return param_temp
+
+ @staticmethod
+ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], ...]:
+ full_param_buffer, *params_split = ctx.saved_tensors
+
+ split_size = full_param_buffer.shape[0] // len(params_split)
+ grads = []
+
+ for i, _ in enumerate(params_split):
+ grads.append(grad_output[i * split_size : (i+1) * split_size])
+
+ return None, *grads
+
+
 class TransformerEngineBaseModule(torch.nn.Module, ABC):
  """Base TE module."""
 
@@ -572,6 +608,29 @@ def grad_output_preprocess(
 
  return grad_output_mat, grad_output_c, grad_output_t, grad_bias
 
+ def noop_cat(self, buffer_name: str, pnames: List[str]) -> torch.Tensor:
+ """No-op replacement of `torch.cat`. The buffer and split parameters must occupy
+ the same memory region. If this is not the case, then the split parameters
+ are concatenated and the buffer is overwritten. The parameters' memory is then
+ re-assigned to point to the buffer to avoid subsequent concatenations.
+ """
+
+ assert hasattr(self, buffer_name), f"No buffer named {buffer_name}"
+ full_param_buffer = getattr(self, buffer_name)
+ split_size = full_param_buffer.shape[0] // len(pnames)
+ params = [getattr(self, name) for name in pnames]
+ for i, p in enumerate(params):
+ if p.data.data_ptr() != full_param_buffer[i*split_size : (i+1)*split_size].data_ptr():
+ with torch.no_grad():
+ setattr(self, buffer_name, torch.cat(params))
+ for j, pname in enumerate(pnames):
+ full_param_buffer = getattr(self, buffer_name)
+ setattr(self, pname,
+ Parameter(full_param_buffer[j*split_size : (j+1)*split_size]))
+ break
+
+ return _NoopCat.apply(getattr(self, buffer_name), *[getattr(self, name) for name in pnames])
+
  @abstractmethod
  def forward(self):
  """Needs override."""
@@ -993,6 +1052,11 @@ class LayerNormLinear(TransformerEngineBaseModule):
  together with the output of the linear transformation.
  Example use case: residual connection for transformer module is
  taken post layernorm.
+ parameters_split : Tuple[str, ...], default = None
+ if a tuple of strings is provided, the weight and bias parameters of the
+ module are exposed as `N` separate `torch.nn.parameter.Parameter`s each,
+ split along the first dimension, where `N` is the length of the argument
+ and the strings contained are the names of the split parameters.
 
  Parallelism parameters
  ----------------------
@@ -1047,6 +1111,7 @@ def __init__(
  parallel_mode: Optional[str] = None,
  return_layernorm_output: bool = False,
  skip_weight_param_allocation: bool = False,
+ parameters_split: Optional[Tuple[str, ...]] = None,
  ) -> None:
  super().__init__()
  self.in_features = in_features
@@ -1055,7 +1120,7 @@ def __init__(
  self.use_bias = bias
  self.return_bias = return_bias
  self.return_layernorm_output = return_layernorm_output
- self.skip_weight_param_allocation = skip_weight_param_allocation
+ self.parameters_split = parameters_split
 
  if tp_group is None:
  self.tp_size = tp_size
@@ -1101,38 +1166,76 @@ def __init__(
  self.reset_layer_norm_parameters()
 
  if not skip_weight_param_allocation:
- self.weight = Parameter(
- torch.empty(
- self.out_features,
- self.in_features,
- device=torch.cuda.current_device(),
- dtype=params_dtype,
- )
- )
+ self.register_buffer("weight_tensor",
+ torch.empty(
+ self.out_features,
+ self.in_features,
+ device=torch.cuda.current_device(),
+ dtype=params_dtype),
+ persistent=False)
 
  initialize_affine_weight_gpu(
- self.weight,
+ self.weight_tensor,
  init_method,
  get_rng_state_tracker,
  partition_dim=1 if self.parallel_mode == "row" else 0,
  stride=1,
  )
 
  if self.use_bias or self.return_bias:
- self.bias = Parameter(
- torch.empty(
- self.out_features,
- device=torch.cuda.current_device(),
- dtype=params_dtype,
- )
- )
- if self.parallel_mode == "column":
- set_tensor_model_parallel_attributes(self.bias, True, 0, 1)
+ self.register_buffer("bias_tensor",
+ torch.empty(
+ self.out_features,
+ device=torch.cuda.current_device(),
+ dtype=params_dtype),
+ persistent=False)
  else:
- self.register_buffer("bias", torch.Tensor().type(params_dtype), persistent=False)
+ self.register_buffer(
+ "bias_tensor", torch.Tensor().type(params_dtype), persistent=False
+ )
 
  with torch.no_grad():
- self.bias.zero_()
+ self.bias_tensor.zero_()
+
+ if parameters_split is None:
+ parameters_split = ("",)
+
+ assert (
+ self.out_features % len(parameters_split) == 0
+ ), f"Weight and bias params cannot be split into {len(parameters_split)} parts"
+
+ split_size = self.out_features // len(parameters_split)
+
+ self.weight_names = []
+ self.bias_names = []
+
+ for i, pname in enumerate(parameters_split):
+ wname = pname + "weight"
+ bname = pname + "bias"
+
+ self.register_parameter(
+ wname, Parameter(self.weight_tensor[i * split_size : (i+1) * split_size])
+ )
+
+ set_tensor_model_parallel_attributes(
+ tensor=getattr(self, wname),
+ is_parallel=True,
+ dim=1 if parallel_mode == "row" else 0,
+ stride=1,
+ )
+
+ if self.use_bias or self.return_bias:
+ self.register_parameter(
+ bname, Parameter(self.bias_tensor[i * split_size : (i+1) * split_size])
+ )
+ else:
+ self.register_buffer(bname, torch.Tensor().type(params_dtype), persistent=False)
+
+ if parallel_mode == "column":
+ set_tensor_model_parallel_attributes(getattr(self, bname), True, 0, 1)
+
+ self.weight_names.append(wname)
+ self.bias_names.append(bname)
 
  self.fp8_weight_shapes.append(torch.Size((self.out_features, self.in_features)))
 
@@ -1193,7 +1296,18 @@ def forward(
  """
 
  with self.prepare_forward(inp, is_first_microbatch) as inp:
- bias_tensor = bias if bias is not None else self.bias
+ bias_tensor = (
+ bias if bias is not None
+ else self.bias if self.parameters_split is None
+ else self.bias_tensor if not self.training
+ else self.noop_cat("bias_tensor", self.bias_names)
+ )
+ weight_tensor = (
+ weight if weight is not None
+ else self.weight if self.parameters_split is None
+ else self.weight_tensor if not self.training
+ else self.noop_cat("weight_tensor", self.weight_names)
+ )
 
  if self.training:
  fwd_fn = _LayerNormLinear.apply
@@ -1205,7 +1319,7 @@ def forward(
  inp,
  self.layer_norm_weight,
  self.layer_norm_bias,
-weight if weight is not None else self.weight,
+weight_tensor,
  self.weight1_fp8 if self.fp8 else None,
  self.weight1_t_fp8 if self.fp8 else None,
  bias_tensor,
@@ -1586,6 +1700,11 @@ class Linear(TransformerEngineBaseModule):
  init_method : Callable, default = `None`
  used for initializing weights in the following way: `init_method(weight)`.
  When set to `None`, defaults to `torch.nn.init.normal_(mean=0.0, std=0.023)`.
+ parameters_split : Tuple[str, ...], default = None
+ if a tuple of strings is provided, the weight and bias parameters of the
+ module are exposed as `N` separate `torch.nn.parameter.Parameter`s each,
+ split along the first dimension, where `N` is the length of the argument
+ and the strings contained are the names of the split parameters.
 
  Parallelism parameters
  ----------------------
@@ -1641,14 +1760,15 @@ def __init__(
  params_dtype: torch.dtype = torch.float32,
  parallel_mode: Optional[str] = None,
  skip_weight_param_allocation: bool = False,
+ parameters_split: Optional[Tuple[str, ...]] = None,
  ) -> None:
  super().__init__()
  self.in_features = in_features
  self.out_features = out_features
  self.fuse_wgrad_accumulation = fuse_wgrad_accumulation
  self.use_bias = bias
  self.return_bias = return_bias
- self.skip_weight_param_allocation = skip_weight_param_allocation
+ self.parameters_split = parameters_split
 
  if tp_group is None:
  self.tp_size = tp_size
@@ -1675,38 +1795,76 @@ def __init__(
  self.sequence_parallel = (self.tp_size > 1) and sequence_parallel
 
  if not skip_weight_param_allocation:
- self.weight = Parameter(
- torch.empty(
- self.out_features,
- self.in_features,
- device=torch.cuda.current_device(),
- dtype=params_dtype,
- )
- )
+ self.register_buffer("weight_tensor",
+ torch.empty(
+ self.out_features,
+ self.in_features,
+ device=torch.cuda.current_device(),
+ dtype=params_dtype),
+ persistent=False)
 
  initialize_affine_weight_gpu(
- self.weight,
+ self.weight_tensor,
  init_method,
  get_rng_state_tracker,
  partition_dim=1 if self.parallel_mode == "row" else 0,
  stride=1,
  )
 
  if self.use_bias or self.return_bias:
- self.bias = Parameter(
- torch.empty(
- self.out_features,
- device=torch.cuda.current_device(),
- dtype=params_dtype,
- )
- )
- if self.parallel_mode == "column":
- set_tensor_model_parallel_attributes(self.bias, True, 0, 1)
+ self.register_buffer("bias_tensor",
+ torch.empty(
+ self.out_features,
+ device=torch.cuda.current_device(),
+ dtype=params_dtype),
+ persistent=False)
  else:
- self.register_buffer("bias", torch.Tensor().type(params_dtype), persistent=False)
+ self.register_buffer(
+ "bias_tensor", torch.Tensor().type(params_dtype), persistent=False
+ )
 
  with torch.no_grad():
- self.bias.zero_()
+ self.bias_tensor.zero_()
+
+ if parameters_split is None:
+ parameters_split = ("",)
+
+ assert (
+ self.out_features % len(parameters_split) == 0
+ ), f"Weight and bias params cannot be split into {len(parameters_split)} parts"
+
+ split_size = self.out_features // len(parameters_split)
+
+ self.weight_names = []
+ self.bias_names = []
+
+ for i, pname in enumerate(parameters_split):
+ wname = pname + "weight"
+ bname = pname + "bias"
+
+ self.register_parameter(
+ wname, Parameter(self.weight_tensor[i * split_size : (i+1) * split_size])
+ )
+
+ set_tensor_model_parallel_attributes(
+ tensor=getattr(self, wname),
+ is_parallel=True,
+ dim=1 if parallel_mode == "row" else 0,
+ stride=1,
+ )
+
+ if self.use_bias or self.return_bias:
+ self.register_parameter(
+ bname, Parameter(self.bias_tensor[i * split_size : (i+1) * split_size])
+ )
+ else:
+ self.register_buffer(bname, torch.Tensor().type(params_dtype), persistent=False)
+
+ if parallel_mode == "column":
+ set_tensor_model_parallel_attributes(getattr(self, bname), True, 0, 1)
+
+ self.weight_names.append(wname)
+ self.bias_names.append(bname)
 
  self.fp8_weight_shapes.append(torch.Size((self.out_features, self.in_features)))
 
@@ -1755,7 +1913,18 @@ def forward(
  """
 
  with self.prepare_forward(inp, is_first_microbatch) as inp:
- bias_tensor = bias if bias is not None else self.bias
+ bias_tensor = (
+ bias if bias is not None
+ else self.bias if self.parameters_split is None
+ else self.bias_tensor if not self.training
+ else self.noop_cat("bias_tensor", self.bias_names)
+ )
+ weight_tensor = (
+ weight if weight is not None
+ else self.weight if self.parameters_split is None
+ else self.weight_tensor if not self.training
+ else self.noop_cat("weight_tensor", self.weight_names)
+ )
 
  if self.training:
  linear_fn = _Linear.apply
@@ -1764,7 +1933,7 @@ def forward(
  linear_fn = _Linear.forward
  args = [None]
  args += (
-weight if weight is not None else self.weight,
+weight_tensor,
  self.weight1_fp8 if self.fp8 else None,
  self.weight1_t_fp8 if self.fp8 else None,
  inp,

transformer_engine/pytorch/module.py

@@ -7,7 +7,7 @@
 import pickle
 import warnings
 from abc import ABC, abstractmethod
-from typing import Union, Optional, Callable, Tuple, Dict, Any, Mapping
+from typing import Union, Optional, Callable, Tuple, Dict, Any, Mapping, List
 from functools import partial
 from contextlib import contextmanager
 
@@ -135,6 +135,42 @@ def _prepare_backward(fp8: bool,
  delete_key_from_amax_buffer(forward=False)
 
 
+class _NoopCat(torch.autograd.Function):
+ """This class is a no-op replacement for `torch.cat`."""
+
+ @staticmethod
+ def forward(ctx,
+ full_param_buffer: torch.Tensor,
+ *params_split: Tuple[torch.Tensor, ...],
+ ) -> torch.Tensor:
+ assert not full_param_buffer.requires_grad, "Buffers should not require gradient"
+ assert (
+ full_param_buffer.shape[0] % len(params_split) == 0
+ ), "Dimensions not compatible for concatenation"
+
+ param_temp = full_param_buffer.new()
+ param_temp.set_(full_param_buffer.storage(),
+ full_param_buffer.storage_offset(),
+ full_param_buffer.size(),
+ full_param_buffer.stride())
+ param_temp.requires_grad = True
+
+ ctx.save_for_backward(full_param_buffer, *params_split)
+ return param_temp
+
+ @staticmethod
+ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], ...]:
+ full_param_buffer, *params_split = ctx.saved_tensors
+
+ split_size = full_param_buffer.shape[0] // len(params_split)
+ grads = []
+
+ for i, _ in enumerate(params_split):
+ grads.append(grad_output[i * split_size : (i+1) * split_size])
+
+ return None, *grads
+
+
 class TransformerEngineBaseModule(torch.nn.Module, ABC):
  """Base TE module."""
 
@@ -572,6 +608,29 @@ def grad_output_preprocess(
 
  return grad_output_mat, grad_output_c, grad_output_t, grad_bias
 
+ def noop_cat(self, buffer_name: str, pnames: List[str]) -> torch.Tensor:
+ """No-op replacement of `torch.cat`. The buffer and split parameters must occupy
+ the same memory region. If this is not the case, then the split parameters
+ are concatenated and the buffer is overwritten. The parameters' memory is then
+ re-assigned to point to the buffer to avoid subsequent concatenations.
+ """
+
+ assert hasattr(self, buffer_name), f"No buffer named {buffer_name}"
+ full_param_buffer = getattr(self, buffer_name)
+ split_size = full_param_buffer.shape[0] // len(pnames)
+ params = [getattr(self, name) for name in pnames]
+ for i, p in enumerate(params):
+ if p.data.data_ptr() != full_param_buffer[i*split_size : (i+1)*split_size].data_ptr():
+ with torch.no_grad():
+ setattr(self, buffer_name, torch.cat(params))
+ for j, pname in enumerate(pnames):
+ full_param_buffer = getattr(self, buffer_name)
+ setattr(self, pname,
+ Parameter(full_param_buffer[j*split_size : (j+1)*split_size]))
+ break
+
+ return _NoopCat.apply(getattr(self, buffer_name), *[getattr(self, name) for name in pnames])
+
  @abstractmethod
  def forward(self):
  """Needs override."""
@@ -993,6 +1052,11 @@ class LayerNormLinear(TransformerEngineBaseModule):
  together with the output of the linear transformation.
  Example use case: residual connection for transformer module is
  taken post layernorm.
+ parameters_split : Tuple[str, ...], default = None
+ if a tuple of strings is provided, the weight and bias parameters of the
+ module are exposed as `N` separate `torch.nn.parameter.Parameter`s each,
+ split along the first dimension, where `N` is the length of the argument
+ and the strings contained are the names of the split parameters.
 
  Parallelism parameters
  ----------------------
@@ -1047,6 +1111,7 @@ def __init__(
  parallel_mode: Optional[str] = None,
  return_layernorm_output: bool = False,
  skip_weight_param_allocation: bool = False,
+ parameters_split: Optional[Tuple[str, ...]] = None,
  ) -> None:
  super().__init__()
  self.in_features = in_features
@@ -1055,7 +1120,7 @@ def __init__(
  self.use_bias = bias
  self.return_bias = return_bias
  self.return_layernorm_output = return_layernorm_output
- self.skip_weight_param_allocation = skip_weight_param_allocation
+ self.parameters_split = parameters_split
 
  if tp_group is None:
  self.tp_size = tp_size
@@ -1101,38 +1166,76 @@ def __init__(
  self.reset_layer_norm_parameters()
 
  if not skip_weight_param_allocation:
- self.weight = Parameter(
- torch.empty(
- self.out_features,
- self.in_features,
- device=torch.cuda.current_device(),
- dtype=params_dtype,
- )
- )
+ self.register_buffer("weight_tensor",
+ torch.empty(
+ self.out_features,
+ self.in_features,
+ device=torch.cuda.current_device(),
+ dtype=params_dtype),
+ persistent=False)
 
  initialize_affine_weight_gpu(
- self.weight,
+ self.weight_tensor,
  init_method,
  get_rng_state_tracker,
  partition_dim=1 if self.parallel_mode == "row" else 0,
  stride=1,
  )
 
  if self.use_bias or self.return_bias:
- self.bias = Parameter(
- torch.empty(
- self.out_features,
- device=torch.cuda.current_device(),
- dtype=params_dtype,
- )
- )
- if self.parallel_mode == "column":
- set_tensor_model_parallel_attributes(self.bias, True, 0, 1)
+ self.register_buffer("bias_tensor",
+ torch.empty(
+ self.out_features,
+ device=torch.cuda.current_device(),
+ dtype=params_dtype),
+ persistent=False)
  else:
- self.register_buffer("bias", torch.Tensor().type(params_dtype), persistent=False)
+ self.register_buffer(
+ "bias_tensor", torch.Tensor().type(params_dtype), persistent=False
+ )
 
  with torch.no_grad():
- self.bias.zero_()
+ self.bias_tensor.zero_()
+
+ if parameters_split is None:
+ parameters_split = ("",)
+
+ assert (
+ self.out_features % len(parameters_split) == 0
+ ), f"Weight and bias params cannot be split into {len(parameters_split)} parts"
+
+ split_size = self.out_features // len(parameters_split)
+
+ self.weight_names = []
+ self.bias_names = []
+
+ for i, pname in enumerate(parameters_split):
+ wname = pname + "weight"
+ bname = pname + "bias"
+
+ self.register_parameter(
+ wname, Parameter(self.weight_tensor[i * split_size : (i+1) * split_size])
+ )
+
+ set_tensor_model_parallel_attributes(
+ tensor=getattr(self, wname),
+ is_parallel=True,
+ dim=1 if parallel_mode == "row" else 0,
+ stride=1,
+ )
+
+ if self.use_bias or self.return_bias:
+ self.register_parameter(
+ bname, Parameter(self.bias_tensor[i * split_size : (i+1) * split_size])
+ )
+ else:
+ self.register_buffer(bname, torch.Tensor().type(params_dtype), persistent=False)
+
+ if parallel_mode == "column":
+ set_tensor_model_parallel_attributes(getattr(self, bname), True, 0, 1)
+
+ self.weight_names.append(wname)
+ self.bias_names.append(bname)
 
  self.fp8_weight_shapes.append(torch.Size((self.out_features, self.in_features)))
 
@@ -1193,7 +1296,18 @@ def forward(
  """
 
  with self.prepare_forward(inp, is_first_microbatch) as inp:
- bias_tensor = bias if bias is not None else self.bias
+ bias_tensor = (
+ bias if bias is not None
+ else self.bias if self.parameters_split is None
+ else self.bias_tensor if not self.training
+ else self.noop_cat("bias_tensor", self.bias_names)
+ )
+ weight_tensor = (
+ weight if weight is not None
+ else self.weight if self.parameters_split is None
+ else self.weight_tensor if not self.training
+ else self.noop_cat("weight_tensor", self.weight_names)
+ )
 
  if self.training:
  fwd_fn = _LayerNormLinear.apply
@@ -1205,7 +1319,7 @@ def forward(
  inp,
  self.layer_norm_weight,
  self.layer_norm_bias,
-weight if weight is not None else self.weight,
+weight_tensor,
  self.weight1_fp8 if self.fp8 else None,
  self.weight1_t_fp8 if self.fp8 else None,
  bias_tensor,
@@ -1586,6 +1700,11 @@ class Linear(TransformerEngineBaseModule):
  init_method : Callable, default = `None`
  used for initializing weights in the following way: `init_method(weight)`.
  When set to `None`, defaults to `torch.nn.init.normal_(mean=0.0, std=0.023)`.
+ parameters_split : Tuple[str, ...], default = None
+ if a tuple of strings is provided, the weight and bias parameters of the
+ module are exposed as `N` separate `torch.nn.parameter.Parameter`s each,
+ split along the first dimension, where `N` is the length of the argument
+ and the strings contained are the names of the split parameters.
 
  Parallelism parameters
  ----------------------
@@ -1641,14 +1760,15 @@ def __init__(
  params_dtype: torch.dtype = torch.float32,
  parallel_mode: Optional[str] = None,
  skip_weight_param_allocation: bool = False,
+ parameters_split: Optional[Tuple[str, ...]] = None,
  ) -> None:
  super().__init__()
  self.in_features = in_features
  self.out_features = out_features
  self.fuse_wgrad_accumulation = fuse_wgrad_accumulation
  self.use_bias = bias
  self.return_bias = return_bias
- self.skip_weight_param_allocation = skip_weight_param_allocation
+ self.parameters_split = parameters_split
 
  if tp_group is None:
  self.tp_size = tp_size
@@ -1675,38 +1795,76 @@ def __init__(
  self.sequence_parallel = (self.tp_size > 1) and sequence_parallel
 
  if not skip_weight_param_allocation:
- self.weight = Parameter(
- torch.empty(
- self.out_features,
- self.in_features,
- device=torch.cuda.current_device(),
- dtype=params_dtype,
- )
- )
+ self.register_buffer("weight_tensor",
+ torch.empty(
+ self.out_features,
+ self.in_features,
+ device=torch.cuda.current_device(),
+ dtype=params_dtype),
+ persistent=False)
 
  initialize_affine_weight_gpu(
- self.weight,
+ self.weight_tensor,
  init_method,
  get_rng_state_tracker,
  partition_dim=1 if self.parallel_mode == "row" else 0,
  stride=1,
  )
 
  if self.use_bias or self.return_bias:
- self.bias = Parameter(
- torch.empty(
- self.out_features,
- device=torch.cuda.current_device(),
- dtype=params_dtype,
- )
- )
- if self.parallel_mode == "column":
- set_tensor_model_parallel_attributes(self.bias, True, 0, 1)
+ self.register_buffer("bias_tensor",
+ torch.empty(
+ self.out_features,
+ device=torch.cuda.current_device(),
+ dtype=params_dtype),
+ persistent=False)
  else:
- self.register_buffer("bias", torch.Tensor().type(params_dtype), persistent=False)
+ self.register_buffer(
+ "bias_tensor", torch.Tensor().type(params_dtype), persistent=False
+ )
 
  with torch.no_grad():
- self.bias.zero_()
+ self.bias_tensor.zero_()
+
+ if parameters_split is None:
+ parameters_split = ("",)
+
+ assert (
+ self.out_features % len(parameters_split) == 0
+ ), f"Weight and bias params cannot be split into {len(parameters_split)} parts"
+
+ split_size = self.out_features // len(parameters_split)
+
+ self.weight_names = []
+ self.bias_names = []
+
+ for i, pname in enumerate(parameters_split):
+ wname = pname + "weight"
+ bname = pname + "bias"
+
+ self.register_parameter(
+ wname, Parameter(self.weight_tensor[i * split_size : (i+1) * split_size])
+ )
+
+ set_tensor_model_parallel_attributes(
+ tensor=getattr(self, wname),
+ is_parallel=True,
+ dim=1 if parallel_mode == "row" else 0,
+ stride=1,
+ )
+
+ if self.use_bias or self.return_bias:
+ self.register_parameter(
+ bname, Parameter(self.bias_tensor[i * split_size : (i+1) * split_size])
+ )
+ else:
+ self.register_buffer(bname, torch.Tensor().type(params_dtype), persistent=False)
+
+ if parallel_mode == "column":
+ set_tensor_model_parallel_attributes(getattr(self, bname), True, 0, 1)
+
+ self.weight_names.append(wname)
+ self.bias_names.append(bname)
 
  self.fp8_weight_shapes.append(torch.Size((self.out_features, self.in_features)))
 
@@ -1755,7 +1913,18 @@ def forward(
  """
 
  with self.prepare_forward(inp, is_first_microbatch) as inp:
- bias_tensor = bias if bias is not None else self.bias
+ bias_tensor = (
+ bias if bias is not None
+ else self.bias if self.parameters_split is None
+ else self.bias_tensor if not self.training
+ else self.noop_cat("bias_tensor", self.bias_names)
+ )
+ weight_tensor = (
+ weight if weight is not None
+ else self.weight if self.parameters_split is None
+ else self.weight_tensor if not self.training
+ else self.noop_cat("weight_tensor", self.weight_names)
+ )
 
  if self.training:
  linear_fn = _Linear.apply
@@ -1764,7 +1933,7 @@ def forward(
  linear_fn = _Linear.forward
  args = [None]
  args += (
-weight if weight is not None else self.weight,
+weight_tensor,
  self.weight1_fp8 if self.fp8 else None,
  self.weight1_t_fp8 if self.fp8 else None,
  inp,

transformer_engine/pytorch/module.py

@@ -7,7 +7,7 @@
 import pickle
 import warnings
 from abc import ABC, abstractmethod
-from typing import Union, Optional, Callable, Tuple, Dict, Any, Mapping
+from typing import Union, Optional, Callable, Tuple, Dict, Any, Mapping, List
 from functools import partial
 from contextlib import contextmanager
 
@@ -135,6 +135,42 @@ def _prepare_backward(fp8: bool,
  delete_key_from_amax_buffer(forward=False)
 
 
+class _NoopCat(torch.autograd.Function):
+ """This class is a no-op replacement for `torch.cat`."""
+
+ @staticmethod
+ def forward(ctx,
+ full_param_buffer: torch.Tensor,
+ *params_split: Tuple[torch.Tensor, ...],
+ ) -> torch.Tensor:
+ assert not full_param_buffer.requires_grad, "Buffers should not require gradient"
+ assert (
+ full_param_buffer.shape[0] % len(params_split) == 0
+ ), "Dimensions not compatible for concatenation"
+
+ param_temp = full_param_buffer.new()
+ param_temp.set_(full_param_buffer.storage(),
+ full_param_buffer.storage_offset(),
+ full_param_buffer.size(),
+ full_param_buffer.stride())
+ param_temp.requires_grad = True
+
+ ctx.save_for_backward(full_param_buffer, *params_split)
+ return param_temp
+
+ @staticmethod
+ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], ...]:
+ full_param_buffer, *params_split = ctx.saved_tensors
+
+ split_size = full_param_buffer.shape[0] // len(params_split)
+ grads = []
+
+ for i, _ in enumerate(params_split):
+ grads.append(grad_output[i * split_size : (i+1) * split_size])
+
+ return None, *grads
+
+
 class TransformerEngineBaseModule(torch.nn.Module, ABC):
  """Base TE module."""
 
@@ -572,6 +608,29 @@ def grad_output_preprocess(
 
  return grad_output_mat, grad_output_c, grad_output_t, grad_bias
 
+ def noop_cat(self, buffer_name: str, pnames: List[str]) -> torch.Tensor:
+ """No-op replacement of `torch.cat`. The buffer and split parameters must occupy
+ the same memory region. If this is not the case, then the split parameters
+ are concatenated and the buffer is overwritten. The parameters' memory is then
+ re-assigned to point to the buffer to avoid subsequent concatenations.
+ """
+
+ assert hasattr(self, buffer_name), f"No buffer named {buffer_name}"
+ full_param_buffer = getattr(self, buffer_name)
+ split_size = full_param_buffer.shape[0] // len(pnames)
+ params = [getattr(self, name) for name in pnames]
+ for i, p in enumerate(params):
+ if p.data.data_ptr() != full_param_buffer[i*split_size : (i+1)*split_size].data_ptr():
+ with torch.no_grad():
+ setattr(self, buffer_name, torch.cat(params))
+ for j, pname in enumerate(pnames):
+ full_param_buffer = getattr(self, buffer_name)
+ setattr(self, pname,
+ Parameter(full_param_buffer[j*split_size : (j+1)*split_size]))
+ break
+
+ return _NoopCat.apply(getattr(self, buffer_name), *[getattr(self, name) for name in pnames])
+
  @abstractmethod
  def forward(self):
  """Needs override."""
@@ -993,6 +1052,11 @@ class LayerNormLinear(TransformerEngineBaseModule):
  together with the output of the linear transformation.
  Example use case: residual connection for transformer module is
  taken post layernorm.
+ parameters_split : Tuple[str, ...], default = None
+ if a tuple of strings is provided, the weight and bias parameters of the
+ module are exposed as `N` separate `torch.nn.parameter.Parameter`s each,
+ split along the first dimension, where `N` is the length of the argument
+ and the strings contained are the names of the split parameters.
 
  Parallelism parameters
  ----------------------
@@ -1047,6 +1111,7 @@ def __init__(
  parallel_mode: Optional[str] = None,
  return_layernorm_output: bool = False,
  skip_weight_param_allocation: bool = False,
+ parameters_split: Optional[Tuple[str, ...]] = None,
  ) -> None:
  super().__init__()
  self.in_features = in_features
@@ -1055,7 +1120,7 @@ def __init__(
  self.use_bias = bias
  self.return_bias = return_bias
  self.return_layernorm_output = return_layernorm_output
- self.skip_weight_param_allocation = skip_weight_param_allocation
+ self.parameters_split = parameters_split
 
  if tp_group is None:
  self.tp_size = tp_size
@@ -1101,38 +1166,76 @@ def __init__(
  self.reset_layer_norm_parameters()
 
  if not skip_weight_param_allocation:
- self.weight = Parameter(
- torch.empty(
- self.out_features,
- self.in_features,
- device=torch.cuda.current_device(),
- dtype=params_dtype,
- )
- )
+ self.register_buffer("weight_tensor",
+ torch.empty(
+ self.out_features,
+ self.in_features,
+ device=torch.cuda.current_device(),
+ dtype=params_dtype),
+ persistent=False)
 
  initialize_affine_weight_gpu(
- self.weight,
+ self.weight_tensor,
  init_method,
  get_rng_state_tracker,
  partition_dim=1 if self.parallel_mode == "row" else 0,
  stride=1,
  )
 
  if self.use_bias or self.return_bias:
- self.bias = Parameter(
- torch.empty(
- self.out_features,
- device=torch.cuda.current_device(),
- dtype=params_dtype,
- )
- )
- if self.parallel_mode == "column":
- set_tensor_model_parallel_attributes(self.bias, True, 0, 1)
+ self.register_buffer("bias_tensor",
+ torch.empty(
+ self.out_features,
+ device=torch.cuda.current_device(),
+ dtype=params_dtype),
+ persistent=False)
  else:
- self.register_buffer("bias", torch.Tensor().type(params_dtype), persistent=False)
+ self.register_buffer(
+ "bias_tensor", torch.Tensor().type(params_dtype), persistent=False
+ )
 
  with torch.no_grad():
- self.bias.zero_()
+ self.bias_tensor.zero_()
+
+ if parameters_split is None:
+ parameters_split = ("",)
+
+ assert (
+ self.out_features % len(parameters_split) == 0
+ ), f"Weight and bias params cannot be split into {len(parameters_split)} parts"
+
+ split_size = self.out_features // len(parameters_split)
+
+ self.weight_names = []
+ self.bias_names = []
+
+ for i, pname in enumerate(parameters_split):
+ wname = pname + "weight"
+ bname = pname + "bias"
+
+ self.register_parameter(
+ wname, Parameter(self.weight_tensor[i * split_size : (i+1) * split_size])
+ )
+
+ set_tensor_model_parallel_attributes(
+ tensor=getattr(self, wname),
+ is_parallel=True,
+ dim=1 if parallel_mode == "row" else 0,
+ stride=1,
+ )
+
+ if self.use_bias or self.return_bias:
+ self.register_parameter(
+ bname, Parameter(self.bias_tensor[i * split_size : (i+1) * split_size])
+ )
+ else:
+ self.register_buffer(bname, torch.Tensor().type(params_dtype), persistent=False)
+
+ if parallel_mode == "column":
+ set_tensor_model_parallel_attributes(getattr(self, bname), True, 0, 1)
+
+ self.weight_names.append(wname)
+ self.bias_names.append(bname)
 
  self.fp8_weight_shapes.append(torch.Size((self.out_features, self.in_features)))
 
@@ -1193,7 +1296,18 @@ def forward(
  """
 
  with self.prepare_forward(inp, is_first_microbatch) as inp:
- bias_tensor = bias if bias is not None else self.bias
+ bias_tensor = (
+ bias if bias is not None
+ else self.bias if self.parameters_split is None
+ else self.bias_tensor if not self.training
+ else self.noop_cat("bias_tensor", self.bias_names)
+ )
+ weight_tensor = (
+ weight if weight is not None
+ else self.weight if self.parameters_split is None
+ else self.weight_tensor if not self.training
+ else self.noop_cat("weight_tensor", self.weight_names)
+ )
 
  if self.training:
  fwd_fn = _LayerNormLinear.apply
@@ -1205,7 +1319,7 @@ def forward(
  inp,
  self.layer_norm_weight,
  self.layer_norm_bias,
-weight if weight is not None else self.weight,
+weight_tensor,
  self.weight1_fp8 if self.fp8 else None,
  self.weight1_t_fp8 if self.fp8 else None,
  bias_tensor,
@@ -1586,6 +1700,11 @@ class Linear(TransformerEngineBaseModule):
  init_method : Callable, default = `None`
  used for initializing weights in the following way: `init_method(weight)`.
  When set to `None`, defaults to `torch.nn.init.normal_(mean=0.0, std=0.023)`.
+ parameters_split : Tuple[str, ...], default = None
+ if a tuple of strings is provided, the weight and bias parameters of the
+ module are exposed as `N` separate `torch.nn.parameter.Parameter`s each,
+ split along the first dimension, where `N` is the length of the argument
+ and the strings contained are the names of the split parameters.
 
  Parallelism parameters
  ----------------------
@@ -1641,14 +1760,15 @@ def __init__(
  params_dtype: torch.dtype = torch.float32,
  parallel_mode: Optional[str] = None,
  skip_weight_param_allocation: bool = False,
+ parameters_split: Optional[Tuple[str, ...]] = None,
  ) -> None:
  super().__init__()
  self.in_features = in_features
  self.out_features = out_features
  self.fuse_wgrad_accumulation = fuse_wgrad_accumulation
  self.use_bias = bias
  self.return_bias = return_bias
- self.skip_weight_param_allocation = skip_weight_param_allocation
+ self.parameters_split = parameters_split
 
  if tp_group is None:
  self.tp_size = tp_size
@@ -1675,38 +1795,76 @@ def __init__(
  self.sequence_parallel = (self.tp_size > 1) and sequence_parallel
 
  if not skip_weight_param_allocation:
- self.weight = Parameter(
- torch.empty(
- self.out_features,
- self.in_features,
- device=torch.cuda.current_device(),
- dtype=params_dtype,
- )
- )
+ self.register_buffer("weight_tensor",
+ torch.empty(
+ self.out_features,
+ self.in_features,
+ device=torch.cuda.current_device(),
+ dtype=params_dtype),
+ persistent=False)
 
  initialize_affine_weight_gpu(
- self.weight,
+ self.weight_tensor,
  init_method,
  get_rng_state_tracker,
  partition_dim=1 if self.parallel_mode == "row" else 0,
  stride=1,
  )
 
  if self.use_bias or self.return_bias:
- self.bias = Parameter(
- torch.empty(
- self.out_features,
- device=torch.cuda.current_device(),
- dtype=params_dtype,
- )
- )
- if self.parallel_mode == "column":
- set_tensor_model_parallel_attributes(self.bias, True, 0, 1)
+ self.register_buffer("bias_tensor",
+ torch.empty(
+ self.out_features,
+ device=torch.cuda.current_device(),
+ dtype=params_dtype),
+ persistent=False)
  else:
- self.register_buffer("bias", torch.Tensor().type(params_dtype), persistent=False)
+ self.register_buffer(
+ "bias_tensor", torch.Tensor().type(params_dtype), persistent=False
+ )
 
  with torch.no_grad():
- self.bias.zero_()
+ self.bias_tensor.zero_()
+
+ if parameters_split is None:
+ parameters_split = ("",)
+
+ assert (
+ self.out_features % len(parameters_split) == 0
+ ), f"Weight and bias params cannot be split into {len(parameters_split)} parts"
+
+ split_size = self.out_features // len(parameters_split)
+
+ self.weight_names = []
+ self.bias_names = []
+
+ for i, pname in enumerate(parameters_split):
+ wname = pname + "weight"
+ bname = pname + "bias"
+
+ self.register_parameter(
+ wname, Parameter(self.weight_tensor[i * split_size : (i+1) * split_size])
+ )
+
+ set_tensor_model_parallel_attributes(
+ tensor=getattr(self, wname),
+ is_parallel=True,
+ dim=1 if parallel_mode == "row" else 0,
+ stride=1,
+ )
+
+ if self.use_bias or self.return_bias:
+ self.register_parameter(
+ bname, Parameter(self.bias_tensor[i * split_size : (i+1) * split_size])
+ )
+ else:
+ self.register_buffer(bname, torch.Tensor().type(params_dtype), persistent=False)
+
+ if parallel_mode == "column":
+ set_tensor_model_parallel_attributes(getattr(self, bname), True, 0, 1)
+
+ self.weight_names.append(wname)
+ self.bias_names.append(bname)
 
  self.fp8_weight_shapes.append(torch.Size((self.out_features, self.in_features)))
 
@@ -1755,7 +1913,18 @@ def forward(
  """
 
  with self.prepare_forward(inp, is_first_microbatch) as inp:
- bias_tensor = bias if bias is not None else self.bias
+ bias_tensor = (
+ bias if bias is not None
+ else self.bias if self.parameters_split is None
+ else self.bias_tensor if not self.training
+ else self.noop_cat("bias_tensor", self.bias_names)
+ )
+ weight_tensor = (
+ weight if weight is not None
+ else self.weight if self.parameters_split is None
+ else self.weight_tensor if not self.training
+ else self.noop_cat("weight_tensor", self.weight_names)
+ )
 
  if self.training:
  linear_fn = _Linear.apply
@@ -1764,7 +1933,7 @@ def forward(
  linear_fn = _Linear.forward
  args = [None]
  args += (
-weight if weight is not None else self.weight,
+weight_tensor,
  self.weight1_fp8 if self.fp8 else None,
  self.weight1_t_fp8 if self.fp8 else None,
  inp,

transformer_engine/pytorch/module.py

@@ -7,7 +7,7 @@
 import pickle
 import warnings
 from abc import ABC, abstractmethod
-from typing import Union, Optional, Callable, Tuple, Dict, Any, Mapping
+from typing import Union, Optional, Callable, Tuple, Dict, Any, Mapping, List
 from functools import partial
 from contextlib import contextmanager
 
@@ -135,6 +135,42 @@ def _prepare_backward(fp8: bool,
  delete_key_from_amax_buffer(forward=False)
 
 
+class _NoopCat(torch.autograd.Function):
+ """This class is a no-op replacement for `torch.cat`."""
+
+ @staticmethod
+ def forward(ctx,
+ full_param_buffer: torch.Tensor,
+ *params_split: Tuple[torch.Tensor, ...],
+ ) -> torch.Tensor:
+ assert not full_param_buffer.requires_grad, "Buffers should not require gradient"
+ assert (
+ full_param_buffer.shape[0] % len(params_split) == 0
+ ), "Dimensions not compatible for concatenation"
+
+ param_temp = full_param_buffer.new()
+ param_temp.set_(full_param_buffer.storage(),
+ full_param_buffer.storage_offset(),
+ full_param_buffer.size(),
+ full_param_buffer.stride())
+ param_temp.requires_grad = True
+
+ ctx.save_for_backward(full_param_buffer, *params_split)
+ return param_temp
+
+ @staticmethod
+ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], ...]:
+ full_param_buffer, *params_split = ctx.saved_tensors
+
+ split_size = full_param_buffer.shape[0] // len(params_split)
+ grads = []
+
+ for i, _ in enumerate(params_split):
+ grads.append(grad_output[i * split_size : (i+1) * split_size])
+
+ return None, *grads
+
+
 class TransformerEngineBaseModule(torch.nn.Module, ABC):
  """Base TE module."""
 
@@ -572,6 +608,29 @@ def grad_output_preprocess(
 
  return grad_output_mat, grad_output_c, grad_output_t, grad_bias
 
+ def noop_cat(self, buffer_name: str, pnames: List[str]) -> torch.Tensor:
+ """No-op replacement of `torch.cat`. The buffer and split parameters must occupy
+ the same memory region. If this is not the case, then the split parameters
+ are concatenated and the buffer is overwritten. The parameters' memory is then
+ re-assigned to point to the buffer to avoid subsequent concatenations.
+ """
+
+ assert hasattr(self, buffer_name), f"No buffer named {buffer_name}"
+ full_param_buffer = getattr(self, buffer_name)
+ split_size = full_param_buffer.shape[0] // len(pnames)
+ params = [getattr(self, name) for name in pnames]
+ for i, p in enumerate(params):
+ if p.data.data_ptr() != full_param_buffer[i*split_size : (i+1)*split_size].data_ptr():
+ with torch.no_grad():
+ setattr(self, buffer_name, torch.cat(params))
+ for j, pname in enumerate(pnames):
+ full_param_buffer = getattr(self, buffer_name)
+ setattr(self, pname,
+ Parameter(full_param_buffer[j*split_size : (j+1)*split_size]))
+ break
+
+ return _NoopCat.apply(getattr(self, buffer_name), *[getattr(self, name) for name in pnames])
+
  @abstractmethod
  def forward(self):
  """Needs override."""
@@ -993,6 +1052,11 @@ class LayerNormLinear(TransformerEngineBaseModule):
  together with the output of the linear transformation.
  Example use case: residual connection for transformer module is
  taken post layernorm.
+ parameters_split : Tuple[str, ...], default = None
+ if a tuple of strings is provided, the weight and bias parameters of the
+ module are exposed as `N` separate `torch.nn.parameter.Parameter`s each,
+ split along the first dimension, where `N` is the length of the argument
+ and the strings contained are the names of the split parameters.
 
  Parallelism parameters
  ----------------------
@@ -1047,6 +1111,7 @@ def __init__(
  parallel_mode: Optional[str] = None,
  return_layernorm_output: bool = False,
  skip_weight_param_allocation: bool = False,
+ parameters_split: Optional[Tuple[str, ...]] = None,
  ) -> None:
  super().__init__()
  self.in_features = in_features
@@ -1055,7 +1120,7 @@ def __init__(
  self.use_bias = bias
  self.return_bias = return_bias
  self.return_layernorm_output = return_layernorm_output
- self.skip_weight_param_allocation = skip_weight_param_allocation
+ self.parameters_split = parameters_split
 
  if tp_group is None:
  self.tp_size = tp_size
@@ -1101,38 +1166,76 @@ def __init__(
  self.reset_layer_norm_parameters()
 
  if not skip_weight_param_allocation:
- self.weight = Parameter(
- torch.empty(
- self.out_features,
- self.in_features,
- device=torch.cuda.current_device(),
- dtype=params_dtype,
- )
- )
+ self.register_buffer("weight_tensor",
+ torch.empty(
+ self.out_features,
+ self.in_features,
+ device=torch.cuda.current_device(),
+ dtype=params_dtype),
+ persistent=False)
 
  initialize_affine_weight_gpu(
- self.weight,
+ self.weight_tensor,
  init_method,
  get_rng_state_tracker,
  partition_dim=1 if self.parallel_mode == "row" else 0,
  stride=1,
  )
 
  if self.use_bias or self.return_bias:
- self.bias = Parameter(
- torch.empty(
- self.out_features,
- device=torch.cuda.current_device(),
- dtype=params_dtype,
- )
- )
- if self.parallel_mode == "column":
- set_tensor_model_parallel_attributes(self.bias, True, 0, 1)
+ self.register_buffer("bias_tensor",
+ torch.empty(
+ self.out_features,
+ device=torch.cuda.current_device(),
+ dtype=params_dtype),
+ persistent=False)
  else:
- self.register_buffer("bias", torch.Tensor().type(params_dtype), persistent=False)
+ self.register_buffer(
+ "bias_tensor", torch.Tensor().type(params_dtype), persistent=False
+ )
 
  with torch.no_grad():
- self.bias.zero_()
+ self.bias_tensor.zero_()
+
+ if parameters_split is None:
+ parameters_split = ("",)
+
+ assert (
+ self.out_features % len(parameters_split) == 0
+ ), f"Weight and bias params cannot be split into {len(parameters_split)} parts"
+
+ split_size = self.out_features // len(parameters_split)
+
+ self.weight_names = []
+ self.bias_names = []
+
+ for i, pname in enumerate(parameters_split):
+ wname = pname + "weight"
+ bname = pname + "bias"
+
+ self.register_parameter(
+ wname, Parameter(self.weight_tensor[i * split_size : (i+1) * split_size])
+ )
+
+ set_tensor_model_parallel_attributes(
+ tensor=getattr(self, wname),
+ is_parallel=True,
+ dim=1 if parallel_mode == "row" else 0,
+ stride=1,
+ )
+
+ if self.use_bias or self.return_bias:
+ self.register_parameter(
+ bname, Parameter(self.bias_tensor[i * split_size : (i+1) * split_size])
+ )
+ else:
+ self.register_buffer(bname, torch.Tensor().type(params_dtype), persistent=False)
+
+ if parallel_mode == "column":
+ set_tensor_model_parallel_attributes(getattr(self, bname), True, 0, 1)
+
+ self.weight_names.append(wname)
+ self.bias_names.append(bname)
 
  self.fp8_weight_shapes.append(torch.Size((self.out_features, self.in_features)))
 
@@ -1193,7 +1296,18 @@ def forward(
  """
 
  with self.prepare_forward(inp, is_first_microbatch) as inp:
- bias_tensor = bias if bias is not None else self.bias
+ bias_tensor = (
+ bias if bias is not None
+ else self.bias if self.parameters_split is None
+ else self.bias_tensor if not self.training
+ else self.noop_cat("bias_tensor", self.bias_names)
+ )
+ weight_tensor = (
+ weight if weight is not None
+ else self.weight if self.parameters_split is None
+ else self.weight_tensor if not self.training
+ else self.noop_cat("weight_tensor", self.weight_names)
+ )
 
  if self.training:
  fwd_fn = _LayerNormLinear.apply
@@ -1205,7 +1319,7 @@ def forward(
  inp,
  self.layer_norm_weight,
  self.layer_norm_bias,
-weight if weight is not None else self.weight,
+weight_tensor,
  self.weight1_fp8 if self.fp8 else None,
  self.weight1_t_fp8 if self.fp8 else None,
  bias_tensor,
@@ -1586,6 +1700,11 @@ class Linear(TransformerEngineBaseModule):
  init_method : Callable, default = `None`
  used for initializing weights in the following way: `init_method(weight)`.
  When set to `None`, defaults to `torch.nn.init.normal_(mean=0.0, std=0.023)`.
+ parameters_split : Tuple[str, ...], default = None
+ if a tuple of strings is provided, the weight and bias parameters of the
+ module are exposed as `N` separate `torch.nn.parameter.Parameter`s each,
+ split along the first dimension, where `N` is the length of the argument
+ and the strings contained are the names of the split parameters.
 
  Parallelism parameters
  ----------------------
@@ -1641,14 +1760,15 @@ def __init__(
  params_dtype: torch.dtype = torch.float32,
  parallel_mode: Optional[str] = None,
  skip_weight_param_allocation: bool = False,
+ parameters_split: Optional[Tuple[str, ...]] = None,
  ) -> None:
  super().__init__()
  self.in_features = in_features
  self.out_features = out_features
  self.fuse_wgrad_accumulation = fuse_wgrad_accumulation
  self.use_bias = bias
  self.return_bias = return_bias
- self.skip_weight_param_allocation = skip_weight_param_allocation
+ self.parameters_split = parameters_split
 
  if tp_group is None:
  self.tp_size = tp_size
@@ -1675,38 +1795,76 @@ def __init__(
  self.sequence_parallel = (self.tp_size > 1) and sequence_parallel
 
  if not skip_weight_param_allocation:
- self.weight = Parameter(
- torch.empty(
- self.out_features,
- self.in_features,
- device=torch.cuda.current_device(),
- dtype=params_dtype,
- )
- )
+ self.register_buffer("weight_tensor",
+ torch.empty(
+ self.out_features,
+ self.in_features,
+ device=torch.cuda.current_device(),
+ dtype=params_dtype),
+ persistent=False)
 
  initialize_affine_weight_gpu(
- self.weight,
+ self.weight_tensor,
  init_method,
  get_rng_state_tracker,
  partition_dim=1 if self.parallel_mode == "row" else 0,
  stride=1,
  )
 
  if self.use_bias or self.return_bias:
- self.bias = Parameter(
- torch.empty(
- self.out_features,
- device=torch.cuda.current_device(),
- dtype=params_dtype,
- )
- )
- if self.parallel_mode == "column":
- set_tensor_model_parallel_attributes(self.bias, True, 0, 1)
+ self.register_buffer("bias_tensor",
+ torch.empty(
+ self.out_features,
+ device=torch.cuda.current_device(),
+ dtype=params_dtype),
+ persistent=False)
  else:
- self.register_buffer("bias", torch.Tensor().type(params_dtype), persistent=False)
+ self.register_buffer(
+ "bias_tensor", torch.Tensor().type(params_dtype), persistent=False
+ )
 
  with torch.no_grad():
- self.bias.zero_()
+ self.bias_tensor.zero_()
+
+ if parameters_split is None:
+ parameters_split = ("",)
+
+ assert (
+ self.out_features % len(parameters_split) == 0
+ ), f"Weight and bias params cannot be split into {len(parameters_split)} parts"
+
+ split_size = self.out_features // len(parameters_split)
+
+ self.weight_names = []
+ self.bias_names = []
+
+ for i, pname in enumerate(parameters_split):
+ wname = pname + "weight"
+ bname = pname + "bias"
+
+ self.register_parameter(
+ wname, Parameter(self.weight_tensor[i * split_size : (i+1) * split_size])
+ )
+
+ set_tensor_model_parallel_attributes(
+ tensor=getattr(self, wname),
+ is_parallel=True,
+ dim=1 if parallel_mode == "row" else 0,
+ stride=1,
+ )
+
+ if self.use_bias or self.return_bias:
+ self.register_parameter(
+ bname, Parameter(self.bias_tensor[i * split_size : (i+1) * split_size])
+ )
+ else:
+ self.register_buffer(bname, torch.Tensor().type(params_dtype), persistent=False)
+
+ if parallel_mode == "column":
+ set_tensor_model_parallel_attributes(getattr(self, bname), True, 0, 1)
+
+ self.weight_names.append(wname)
+ self.bias_names.append(bname)
 
  self.fp8_weight_shapes.append(torch.Size((self.out_features, self.in_features)))
 
@@ -1755,7 +1913,18 @@ def forward(
  """
 
  with self.prepare_forward(inp, is_first_microbatch) as inp:
- bias_tensor = bias if bias is not None else self.bias
+ bias_tensor = (
+ bias if bias is not None
+ else self.bias if self.parameters_split is None
+ else self.bias_tensor if not self.training
+ else self.noop_cat("bias_tensor", self.bias_names)
+ )
+ weight_tensor = (
+ weight if weight is not None
+ else self.weight if self.parameters_split is None
+ else self.weight_tensor if not self.training
+ else self.noop_cat("weight_tensor", self.weight_names)
+ )
 
  if self.training:
  linear_fn = _Linear.apply
@@ -1764,7 +1933,7 @@ def forward(
  linear_fn = _Linear.forward
  args = [None]
  args += (
-weight if weight is not None else self.weight,
+weight_tensor,
  self.weight1_fp8 if self.fp8 else None,
  self.weight1_t_fp8 if self.fp8 else None,
  inp,

transformer_engine/pytorch/module.py

@@ -7,7 +7,7 @@
 import pickle
 import warnings
 from abc import ABC, abstractmethod
-from typing import Union, Optional, Callable, Tuple, Dict, Any, Mapping
+from typing import Union, Optional, Callable, Tuple, Dict, Any, Mapping, List
 from functools import partial
 from contextlib import contextmanager
 
@@ -135,6 +135,42 @@ def _prepare_backward(fp8: bool,
  delete_key_from_amax_buffer(forward=False)
 
 
+class _NoopCat(torch.autograd.Function):
+ """This class is a no-op replacement for `torch.cat`."""
+
+ @staticmethod
+ def forward(ctx,
+ full_param_buffer: torch.Tensor,
+ *params_split: Tuple[torch.Tensor, ...],
+ ) -> torch.Tensor:
+ assert not full_param_buffer.requires_grad, "Buffers should not require gradient"
+ assert (
+ full_param_buffer.shape[0] % len(params_split) == 0
+ ), "Dimensions not compatible for concatenation"
+
+ param_temp = full_param_buffer.new()
+ param_temp.set_(full_param_buffer.storage(),
+ full_param_buffer.storage_offset(),
+ full_param_buffer.size(),
+ full_param_buffer.stride())
+ param_temp.requires_grad = True
+
+ ctx.save_for_backward(full_param_buffer, *params_split)
+ return param_temp
+
+ @staticmethod
+ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], ...]:
+ full_param_buffer, *params_split = ctx.saved_tensors
+
+ split_size = full_param_buffer.shape[0] // len(params_split)
+ grads = []
+
+ for i, _ in enumerate(params_split):
+ grads.append(grad_output[i * split_size : (i+1) * split_size])
+
+ return None, *grads
+
+
 class TransformerEngineBaseModule(torch.nn.Module, ABC):
  """Base TE module."""
 
@@ -572,6 +608,29 @@ def grad_output_preprocess(
 
  return grad_output_mat, grad_output_c, grad_output_t, grad_bias
 
+ def noop_cat(self, buffer_name: str, pnames: List[str]) -> torch.Tensor:
+ """No-op replacement of `torch.cat`. The buffer and split parameters must occupy
+ the same memory region. If this is not the case, then the split parameters
+ are concatenated and the buffer is overwritten. The parameters' memory is then
+ re-assigned to point to the buffer to avoid subsequent concatenations.
+ """
+
+ assert hasattr(self, buffer_name), f"No buffer named {buffer_name}"
+ full_param_buffer = getattr(self, buffer_name)
+ split_size = full_param_buffer.shape[0] // len(pnames)
+ params = [getattr(self, name) for name in pnames]
+ for i, p in enumerate(params):
+ if p.data.data_ptr() != full_param_buffer[i*split_size : (i+1)*split_size].data_ptr():
+ with torch.no_grad():
+ setattr(self, buffer_name, torch.cat(params))
+ for j, pname in enumerate(pnames):
+ full_param_buffer = getattr(self, buffer_name)
+ setattr(self, pname,
+ Parameter(full_param_buffer[j*split_size : (j+1)*split_size]))
+ break
+
+ return _NoopCat.apply(getattr(self, buffer_name), *[getattr(self, name) for name in pnames])
+
  @abstractmethod
  def forward(self):
  """Needs override."""
@@ -993,6 +1052,11 @@ class LayerNormLinear(TransformerEngineBaseModule):
  together with the output of the linear transformation.
  Example use case: residual connection for transformer module is
  taken post layernorm.
+ parameters_split : Tuple[str, ...], default = None
+ if a tuple of strings is provided, the weight and bias parameters of the
+ module are exposed as `N` separate `torch.nn.parameter.Parameter`s each,
+ split along the first dimension, where `N` is the length of the argument
+ and the strings contained are the names of the split parameters.
 
  Parallelism parameters
  ----------------------
@@ -1047,6 +1111,7 @@ def __init__(
  parallel_mode: Optional[str] = None,
  return_layernorm_output: bool = False,
  skip_weight_param_allocation: bool = False,
+ parameters_split: Optional[Tuple[str, ...]] = None,
  ) -> None:
  super().__init__()
  self.in_features = in_features
@@ -1055,7 +1120,7 @@ def __init__(
  self.use_bias = bias
  self.return_bias = return_bias
  self.return_layernorm_output = return_layernorm_output
- self.skip_weight_param_allocation = skip_weight_param_allocation
+ self.parameters_split = parameters_split
 
  if tp_group is None:
  self.tp_size = tp_size
@@ -1101,38 +1166,76 @@ def __init__(
  self.reset_layer_norm_parameters()
 
  if not skip_weight_param_allocation:
- self.weight = Parameter(
- torch.empty(
- self.out_features,
- self.in_features,
- device=torch.cuda.current_device(),
- dtype=params_dtype,
- )
- )
+ self.register_buffer("weight_tensor",
+ torch.empty(
+ self.out_features,
+ self.in_features,
+ device=torch.cuda.current_device(),
+ dtype=params_dtype),
+ persistent=False)
 
  initialize_affine_weight_gpu(
- self.weight,
+ self.weight_tensor,
  init_method,
  get_rng_state_tracker,
  partition_dim=1 if self.parallel_mode == "row" else 0,
  stride=1,
  )
 
  if self.use_bias or self.return_bias:
- self.bias = Parameter(
- torch.empty(
- self.out_features,
- device=torch.cuda.current_device(),
- dtype=params_dtype,
- )
- )
- if self.parallel_mode == "column":
- set_tensor_model_parallel_attributes(self.bias, True, 0, 1)
+ self.register_buffer("bias_tensor",
+ torch.empty(
+ self.out_features,
+ device=torch.cuda.current_device(),
+ dtype=params_dtype),
+ persistent=False)
  else:
- self.register_buffer("bias", torch.Tensor().type(params_dtype), persistent=False)
+ self.register_buffer(
+ "bias_tensor", torch.Tensor().type(params_dtype), persistent=False
+ )
 
  with torch.no_grad():
- self.bias.zero_()
+ self.bias_tensor.zero_()
+
+ if parameters_split is None:
+ parameters_split = ("",)
+
+ assert (
+ self.out_features % len(parameters_split) == 0
+ ), f"Weight and bias params cannot be split into {len(parameters_split)} parts"
+
+ split_size = self.out_features // len(parameters_split)
+
+ self.weight_names = []
+ self.bias_names = []
+
+ for i, pname in enumerate(parameters_split):
+ wname = pname + "weight"
+ bname = pname + "bias"
+
+ self.register_parameter(
+ wname, Parameter(self.weight_tensor[i * split_size : (i+1) * split_size])
+ )
+
+ set_tensor_model_parallel_attributes(
+ tensor=getattr(self, wname),
+ is_parallel=True,
+ dim=1 if parallel_mode == "row" else 0,
+ stride=1,
+ )
+
+ if self.use_bias or self.return_bias:
+ self.register_parameter(
+ bname, Parameter(self.bias_tensor[i * split_size : (i+1) * split_size])
+ )
+ else:
+ self.register_buffer(bname, torch.Tensor().type(params_dtype), persistent=False)
+
+ if parallel_mode == "column":
+ set_tensor_model_parallel_attributes(getattr(self, bname), True, 0, 1)
+
+ self.weight_names.append(wname)
+ self.bias_names.append(bname)
 
  self.fp8_weight_shapes.append(torch.Size((self.out_features, self.in_features)))
 
@@ -1193,7 +1296,18 @@ def forward(
  """
 
  with self.prepare_forward(inp, is_first_microbatch) as inp:
- bias_tensor = bias if bias is not None else self.bias
+ bias_tensor = (
+ bias if bias is not None
+ else self.bias if self.parameters_split is None
+ else self.bias_tensor if not self.training
+ else self.noop_cat("bias_tensor", self.bias_names)
+ )
+ weight_tensor = (
+ weight if weight is not None
+ else self.weight if self.parameters_split is None
+ else self.weight_tensor if not self.training
+ else self.noop_cat("weight_tensor", self.weight_names)
+ )
 
  if self.training:
  fwd_fn = _LayerNormLinear.apply
@@ -1205,7 +1319,7 @@ def forward(
  inp,
  self.layer_norm_weight,
  self.layer_norm_bias,
-weight if weight is not None else self.weight,
+weight_tensor,
  self.weight1_fp8 if self.fp8 else None,
  self.weight1_t_fp8 if self.fp8 else None,
  bias_tensor,
@@ -1586,6 +1700,11 @@ class Linear(TransformerEngineBaseModule):
  init_method : Callable, default = `None`
  used for initializing weights in the following way: `init_method(weight)`.
  When set to `None`, defaults to `torch.nn.init.normal_(mean=0.0, std=0.023)`.
+ parameters_split : Tuple[str, ...], default = None
+ if a tuple of strings is provided, the weight and bias parameters of the
+ module are exposed as `N` separate `torch.nn.parameter.Parameter`s each,
+ split along the first dimension, where `N` is the length of the argument
+ and the strings contained are the names of the split parameters.
 
  Parallelism parameters
  ----------------------
@@ -1641,14 +1760,15 @@ def __init__(
  params_dtype: torch.dtype = torch.float32,
  parallel_mode: Optional[str] = None,
  skip_weight_param_allocation: bool = False,
+ parameters_split: Optional[Tuple[str, ...]] = None,
  ) -> None:
  super().__init__()
  self.in_features = in_features
  self.out_features = out_features
  self.fuse_wgrad_accumulation = fuse_wgrad_accumulation
  self.use_bias = bias
  self.return_bias = return_bias
- self.skip_weight_param_allocation = skip_weight_param_allocation
+ self.parameters_split = parameters_split
 
  if tp_group is None:
  self.tp_size = tp_size
@@ -1675,38 +1795,76 @@ def __init__(
  self.sequence_parallel = (self.tp_size > 1) and sequence_parallel
 
  if not skip_weight_param_allocation:
- self.weight = Parameter(
- torch.empty(
- self.out_features,
- self.in_features,
- device=torch.cuda.current_device(),
- dtype=params_dtype,
- )
- )
+ self.register_buffer("weight_tensor",
+ torch.empty(
+ self.out_features,
+ self.in_features,
+ device=torch.cuda.current_device(),
+ dtype=params_dtype),
+ persistent=False)
 
  initialize_affine_weight_gpu(
- self.weight,
+ self.weight_tensor,
  init_method,
  get_rng_state_tracker,
  partition_dim=1 if self.parallel_mode == "row" else 0,
  stride=1,
  )
 
  if self.use_bias or self.return_bias:
- self.bias = Parameter(
- torch.empty(
- self.out_features,
- device=torch.cuda.current_device(),
- dtype=params_dtype,
- )
- )
- if self.parallel_mode == "column":
- set_tensor_model_parallel_attributes(self.bias, True, 0, 1)
+ self.register_buffer("bias_tensor",
+ torch.empty(
+ self.out_features,
+ device=torch.cuda.current_device(),
+ dtype=params_dtype),
+ persistent=False)
  else:
- self.register_buffer("bias", torch.Tensor().type(params_dtype), persistent=False)
+ self.register_buffer(
+ "bias_tensor", torch.Tensor().type(params_dtype), persistent=False
+ )
 
  with torch.no_grad():
- self.bias.zero_()
+ self.bias_tensor.zero_()
+
+ if parameters_split is None:
+ parameters_split = ("",)
+
+ assert (
+ self.out_features % len(parameters_split) == 0
+ ), f"Weight and bias params cannot be split into {len(parameters_split)} parts"
+
+ split_size = self.out_features // len(parameters_split)
+
+ self.weight_names = []
+ self.bias_names = []
+
+ for i, pname in enumerate(parameters_split):
+ wname = pname + "weight"
+ bname = pname + "bias"
+
+ self.register_parameter(
+ wname, Parameter(self.weight_tensor[i * split_size : (i+1) * split_size])
+ )
+
+ set_tensor_model_parallel_attributes(
+ tensor=getattr(self, wname),
+ is_parallel=True,
+ dim=1 if parallel_mode == "row" else 0,
+ stride=1,
+ )
+
+ if self.use_bias or self.return_bias:
+ self.register_parameter(
+ bname, Parameter(self.bias_tensor[i * split_size : (i+1) * split_size])
+ )
+ else:
+ self.register_buffer(bname, torch.Tensor().type(params_dtype), persistent=False)
+
+ if parallel_mode == "column":
+ set_tensor_model_parallel_attributes(getattr(self, bname), True, 0, 1)
+
+ self.weight_names.append(wname)
+ self.bias_names.append(bname)
 
  self.fp8_weight_shapes.append(torch.Size((self.out_features, self.in_features)))
 
@@ -1755,7 +1913,18 @@ def forward(
  """
 
  with self.prepare_forward(inp, is_first_microbatch) as inp:
- bias_tensor = bias if bias is not None else self.bias
+ bias_tensor = (
+ bias if bias is not None
+ else self.bias if self.parameters_split is None
+ else self.bias_tensor if not self.training
+ else self.noop_cat("bias_tensor", self.bias_names)
+ )
+ weight_tensor = (
+ weight if weight is not None
+ else self.weight if self.parameters_split is None
+ else self.weight_tensor if not self.training
+ else self.noop_cat("weight_tensor", self.weight_names)
+ )
 
  if self.training:
  linear_fn = _Linear.apply
@@ -1764,7 +1933,7 @@ def forward(
  linear_fn = _Linear.forward
  args = [None]
  args += (
-weight if weight is not None else self.weight,
+weight_tensor,
  self.weight1_fp8 if self.fp8 else None,
  self.weight1_t_fp8 if self.fp8 else None,
  inp,

transformer_engine/pytorch/module.py

@@ -7,7 +7,7 @@
 import pickle
 import warnings
 from abc import ABC, abstractmethod
-from typing import Union, Optional, Callable, Tuple, Dict, Any, Mapping
+from typing import Union, Optional, Callable, Tuple, Dict, Any, Mapping, List
 from functools import partial
 from contextlib import contextmanager
 
@@ -135,6 +135,42 @@ def _prepare_backward(fp8: bool,
  delete_key_from_amax_buffer(forward=False)
 
 
+class _NoopCat(torch.autograd.Function):
+ """This class is a no-op replacement for `torch.cat`."""
+
+ @staticmethod
+ def forward(ctx,
+ full_param_buffer: torch.Tensor,
+ *params_split: Tuple[torch.Tensor, ...],
+ ) -> torch.Tensor:
+ assert not full_param_buffer.requires_grad, "Buffers should not require gradient"
+ assert (
+ full_param_buffer.shape[0] % len(params_split) == 0
+ ), "Dimensions not compatible for concatenation"
+
+ param_temp = full_param_buffer.new()
+ param_temp.set_(full_param_buffer.storage(),
+ full_param_buffer.storage_offset(),
+ full_param_buffer.size(),
+ full_param_buffer.stride())
+ param_temp.requires_grad = True
+
+ ctx.save_for_backward(full_param_buffer, *params_split)
+ return param_temp
+
+ @staticmethod
+ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], ...]:
+ full_param_buffer, *params_split = ctx.saved_tensors
+
+ split_size = full_param_buffer.shape[0] // len(params_split)
+ grads = []
+
+ for i, _ in enumerate(params_split):
+ grads.append(grad_output[i * split_size : (i+1) * split_size])
+
+ return None, *grads
+
+
 class TransformerEngineBaseModule(torch.nn.Module, ABC):
  """Base TE module."""
 
@@ -572,6 +608,29 @@ def grad_output_preprocess(
 
  return grad_output_mat, grad_output_c, grad_output_t, grad_bias
 
+ def noop_cat(self, buffer_name: str, pnames: List[str]) -> torch.Tensor:
+ """No-op replacement of `torch.cat`. The buffer and split parameters must occupy
+ the same memory region. If this is not the case, then the split parameters
+ are concatenated and the buffer is overwritten. The parameters' memory is then
+ re-assigned to point to the buffer to avoid subsequent concatenations.
+ """
+
+ assert hasattr(self, buffer_name), f"No buffer named {buffer_name}"
+ full_param_buffer = getattr(self, buffer_name)
+ split_size = full_param_buffer.shape[0] // len(pnames)
+ params = [getattr(self, name) for name in pnames]
+ for i, p in enumerate(params):
+ if p.data.data_ptr() != full_param_buffer[i*split_size : (i+1)*split_size].data_ptr():
+ with torch.no_grad():
+ setattr(self, buffer_name, torch.cat(params))
+ for j, pname in enumerate(pnames):
+ full_param_buffer = getattr(self, buffer_name)
+ setattr(self, pname,
+ Parameter(full_param_buffer[j*split_size : (j+1)*split_size]))
+ break
+
+ return _NoopCat.apply(getattr(self, buffer_name), *[getattr(self, name) for name in pnames])
+
  @abstractmethod
  def forward(self):
  """Needs override."""
@@ -993,6 +1052,11 @@ class LayerNormLinear(TransformerEngineBaseModule):
  together with the output of the linear transformation.
  Example use case: residual connection for transformer module is
  taken post layernorm.
+ parameters_split : Tuple[str, ...], default = None
+ if a tuple of strings is provided, the weight and bias parameters of the
+ module are exposed as `N` separate `torch.nn.parameter.Parameter`s each,
+ split along the first dimension, where `N` is the length of the argument
+ and the strings contained are the names of the split parameters.
 
  Parallelism parameters
  ----------------------
@@ -1047,6 +1111,7 @@ def __init__(
  parallel_mode: Optional[str] = None,
  return_layernorm_output: bool = False,
  skip_weight_param_allocation: bool = False,
+ parameters_split: Optional[Tuple[str, ...]] = None,
  ) -> None:
  super().__init__()
  self.in_features = in_features
@@ -1055,7 +1120,7 @@ def __init__(
  self.use_bias = bias
  self.return_bias = return_bias
  self.return_layernorm_output = return_layernorm_output
- self.skip_weight_param_allocation = skip_weight_param_allocation
+ self.parameters_split = parameters_split
 
  if tp_group is None:
  self.tp_size = tp_size
@@ -1101,38 +1166,76 @@ def __init__(
  self.reset_layer_norm_parameters()
 
  if not skip_weight_param_allocation:
- self.weight = Parameter(
- torch.empty(
- self.out_features,
- self.in_features,
- device=torch.cuda.current_device(),
- dtype=params_dtype,
- )
- )
+ self.register_buffer("weight_tensor",
+ torch.empty(
+ self.out_features,
+ self.in_features,
+ device=torch.cuda.current_device(),
+ dtype=params_dtype),
+ persistent=False)
 
  initialize_affine_weight_gpu(
- self.weight,
+ self.weight_tensor,
  init_method,
  get_rng_state_tracker,
  partition_dim=1 if self.parallel_mode == "row" else 0,
  stride=1,
  )
 
  if self.use_bias or self.return_bias:
- self.bias = Parameter(
- torch.empty(
- self.out_features,
- device=torch.cuda.current_device(),
- dtype=params_dtype,
- )
- )
- if self.parallel_mode == "column":
- set_tensor_model_parallel_attributes(self.bias, True, 0, 1)
+ self.register_buffer("bias_tensor",
+ torch.empty(
+ self.out_features,
+ device=torch.cuda.current_device(),
+ dtype=params_dtype),
+ persistent=False)
  else:
- self.register_buffer("bias", torch.Tensor().type(params_dtype), persistent=False)
+ self.register_buffer(
+ "bias_tensor", torch.Tensor().type(params_dtype), persistent=False
+ )
 
  with torch.no_grad():
- self.bias.zero_()
+ self.bias_tensor.zero_()
+
+ if parameters_split is None:
+ parameters_split = ("",)
+
+ assert (
+ self.out_features % len(parameters_split) == 0
+ ), f"Weight and bias params cannot be split into {len(parameters_split)} parts"
+
+ split_size = self.out_features // len(parameters_split)
+
+ self.weight_names = []
+ self.bias_names = []
+
+ for i, pname in enumerate(parameters_split):
+ wname = pname + "weight"
+ bname = pname + "bias"
+
+ self.register_parameter(
+ wname, Parameter(self.weight_tensor[i * split_size : (i+1) * split_size])
+ )
+
+ set_tensor_model_parallel_attributes(
+ tensor=getattr(self, wname),
+ is_parallel=True,
+ dim=1 if parallel_mode == "row" else 0,
+ stride=1,
+ )
+
+ if self.use_bias or self.return_bias:
+ self.register_parameter(
+ bname, Parameter(self.bias_tensor[i * split_size : (i+1) * split_size])
+ )
+ else:
+ self.register_buffer(bname, torch.Tensor().type(params_dtype), persistent=False)
+
+ if parallel_mode == "column":
+ set_tensor_model_parallel_attributes(getattr(self, bname), True, 0, 1)
+
+ self.weight_names.append(wname)
+ self.bias_names.append(bname)
 
  self.fp8_weight_shapes.append(torch.Size((self.out_features, self.in_features)))
 
@@ -1193,7 +1296,18 @@ def forward(
  """
 
  with self.prepare_forward(inp, is_first_microbatch) as inp:
- bias_tensor = bias if bias is not None else self.bias
+ bias_tensor = (
+ bias if bias is not None
+ else self.bias if self.parameters_split is None
+ else self.bias_tensor if not self.training
+ else self.noop_cat("bias_tensor", self.bias_names)
+ )
+ weight_tensor = (
+ weight if weight is not None
+ else self.weight if self.parameters_split is None
+ else self.weight_tensor if not self.training
+ else self.noop_cat("weight_tensor", self.weight_names)
+ )
 
  if self.training:
  fwd_fn = _LayerNormLinear.apply
@@ -1205,7 +1319,7 @@ def forward(
  inp,
  self.layer_norm_weight,
  self.layer_norm_bias,
-weight if weight is not None else self.weight,
+weight_tensor,
  self.weight1_fp8 if self.fp8 else None,
  self.weight1_t_fp8 if self.fp8 else None,
  bias_tensor,
@@ -1586,6 +1700,11 @@ class Linear(TransformerEngineBaseModule):
  init_method : Callable, default = `None`
  used for initializing weights in the following way: `init_method(weight)`.
  When set to `None`, defaults to `torch.nn.init.normal_(mean=0.0, std=0.023)`.
+ parameters_split : Tuple[str, ...], default = None
+ if a tuple of strings is provided, the weight and bias parameters of the
+ module are exposed as `N` separate `torch.nn.parameter.Parameter`s each,
+ split along the first dimension, where `N` is the length of the argument
+ and the strings contained are the names of the split parameters.
 
  Parallelism parameters
  ----------------------
@@ -1641,14 +1760,15 @@ def __init__(
  params_dtype: torch.dtype = torch.float32,
  parallel_mode: Optional[str] = None,
  skip_weight_param_allocation: bool = False,
+ parameters_split: Optional[Tuple[str, ...]] = None,
  ) -> None:
  super().__init__()
  self.in_features = in_features
  self.out_features = out_features
  self.fuse_wgrad_accumulation = fuse_wgrad_accumulation
  self.use_bias = bias
  self.return_bias = return_bias
- self.skip_weight_param_allocation = skip_weight_param_allocation
+ self.parameters_split = parameters_split
 
  if tp_group is None:
  self.tp_size = tp_size
@@ -1675,38 +1795,76 @@ def __init__(
  self.sequence_parallel = (self.tp_size > 1) and sequence_parallel
 
  if not skip_weight_param_allocation:
- self.weight = Parameter(
- torch.empty(
- self.out_features,
- self.in_features,
- device=torch.cuda.current_device(),
- dtype=params_dtype,
- )
- )
+ self.register_buffer("weight_tensor",
+ torch.empty(
+ self.out_features,
+ self.in_features,
+ device=torch.cuda.current_device(),
+ dtype=params_dtype),
+ persistent=False)
 
  initialize_affine_weight_gpu(
- self.weight,
+ self.weight_tensor,
  init_method,
  get_rng_state_tracker,
  partition_dim=1 if self.parallel_mode == "row" else 0,
  stride=1,
  )
 
  if self.use_bias or self.return_bias:
- self.bias = Parameter(
- torch.empty(
- self.out_features,
- device=torch.cuda.current_device(),
- dtype=params_dtype,
- )
- )
- if self.parallel_mode == "column":
- set_tensor_model_parallel_attributes(self.bias, True, 0, 1)
+ self.register_buffer("bias_tensor",
+ torch.empty(
+ self.out_features,
+ device=torch.cuda.current_device(),
+ dtype=params_dtype),
+ persistent=False)
  else:
- self.register_buffer("bias", torch.Tensor().type(params_dtype), persistent=False)
+ self.register_buffer(
+ "bias_tensor", torch.Tensor().type(params_dtype), persistent=False
+ )
 
  with torch.no_grad():
- self.bias.zero_()
+ self.bias_tensor.zero_()
+
+ if parameters_split is None:
+ parameters_split = ("",)
+
+ assert (
+ self.out_features % len(parameters_split) == 0
+ ), f"Weight and bias params cannot be split into {len(parameters_split)} parts"
+
+ split_size = self.out_features // len(parameters_split)
+
+ self.weight_names = []
+ self.bias_names = []
+
+ for i, pname in enumerate(parameters_split):
+ wname = pname + "weight"
+ bname = pname + "bias"
+
+ self.register_parameter(
+ wname, Parameter(self.weight_tensor[i * split_size : (i+1) * split_size])
+ )
+
+ set_tensor_model_parallel_attributes(
+ tensor=getattr(self, wname),
+ is_parallel=True,
+ dim=1 if parallel_mode == "row" else 0,
+ stride=1,
+ )
+
+ if self.use_bias or self.return_bias:
+ self.register_parameter(
+ bname, Parameter(self.bias_tensor[i * split_size : (i+1) * split_size])
+ )
+ else:
+ self.register_buffer(bname, torch.Tensor().type(params_dtype), persistent=False)
+
+ if parallel_mode == "column":
+ set_tensor_model_parallel_attributes(getattr(self, bname), True, 0, 1)
+
+ self.weight_names.append(wname)
+ self.bias_names.append(bname)
 
  self.fp8_weight_shapes.append(torch.Size((self.out_features, self.in_features)))
 
@@ -1755,7 +1913,18 @@ def forward(
  """
 
  with self.prepare_forward(inp, is_first_microbatch) as inp:
- bias_tensor = bias if bias is not None else self.bias
+ bias_tensor = (
+ bias if bias is not None
+ else self.bias if self.parameters_split is None
+ else self.bias_tensor if not self.training
+ else self.noop_cat("bias_tensor", self.bias_names)
+ )
+ weight_tensor = (
+ weight if weight is not None
+ else self.weight if self.parameters_split is None
+ else self.weight_tensor if not self.training
+ else self.noop_cat("weight_tensor", self.weight_names)
+ )
 
  if self.training:
  linear_fn = _Linear.apply
@@ -1764,7 +1933,7 @@ def forward(
  linear_fn = _Linear.forward
  args = [None]
  args += (
-weight if weight is not None else self.weight,
+weight_tensor,
  self.weight1_fp8 if self.fp8 else None,
  self.weight1_t_fp8 if self.fp8 else None,
  inp,

transformer_engine/pytorch/module.py

@@ -7,7 +7,7 @@
 import pickle
 import warnings
 from abc import ABC, abstractmethod
-from typing import Union, Optional, Callable, Tuple, Dict, Any, Mapping
+from typing import Union, Optional, Callable, Tuple, Dict, Any, Mapping, List
 from functools import partial
 from contextlib import contextmanager
 
@@ -135,6 +135,42 @@ def _prepare_backward(fp8: bool,
  delete_key_from_amax_buffer(forward=False)
 
 
+class _NoopCat(torch.autograd.Function):
+ """This class is a no-op replacement for `torch.cat`."""
+
+ @staticmethod
+ def forward(ctx,
+ full_param_buffer: torch.Tensor,
+ *params_split: Tuple[torch.Tensor, ...],
+ ) -> torch.Tensor:
+ assert not full_param_buffer.requires_grad, "Buffers should not require gradient"
+ assert (
+ full_param_buffer.shape[0] % len(params_split) == 0
+ ), "Dimensions not compatible for concatenation"
+
+ param_temp = full_param_buffer.new()
+ param_temp.set_(full_param_buffer.storage(),
+ full_param_buffer.storage_offset(),
+ full_param_buffer.size(),
+ full_param_buffer.stride())
+ param_temp.requires_grad = True
+
+ ctx.save_for_backward(full_param_buffer, *params_split)
+ return param_temp
+
+ @staticmethod
+ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], ...]:
+ full_param_buffer, *params_split = ctx.saved_tensors
+
+ split_size = full_param_buffer.shape[0] // len(params_split)
+ grads = []
+
+ for i, _ in enumerate(params_split):
+ grads.append(grad_output[i * split_size : (i+1) * split_size])
+
+ return None, *grads
+
+
 class TransformerEngineBaseModule(torch.nn.Module, ABC):
  """Base TE module."""
 
@@ -572,6 +608,29 @@ def grad_output_preprocess(
 
  return grad_output_mat, grad_output_c, grad_output_t, grad_bias
 
+ def noop_cat(self, buffer_name: str, pnames: List[str]) -> torch.Tensor:
+ """No-op replacement of `torch.cat`. The buffer and split parameters must occupy
+ the same memory region. If this is not the case, then the split parameters
+ are concatenated and the buffer is overwritten. The parameters' memory is then
+ re-assigned to point to the buffer to avoid subsequent concatenations.
+ """
+
+ assert hasattr(self, buffer_name), f"No buffer named {buffer_name}"
+ full_param_buffer = getattr(self, buffer_name)
+ split_size = full_param_buffer.shape[0] // len(pnames)
+ params = [getattr(self, name) for name in pnames]
+ for i, p in enumerate(params):
+ if p.data.data_ptr() != full_param_buffer[i*split_size : (i+1)*split_size].data_ptr():
+ with torch.no_grad():
+ setattr(self, buffer_name, torch.cat(params))
+ for j, pname in enumerate(pnames):
+ full_param_buffer = getattr(self, buffer_name)
+ setattr(self, pname,
+ Parameter(full_param_buffer[j*split_size : (j+1)*split_size]))
+ break
+
+ return _NoopCat.apply(getattr(self, buffer_name), *[getattr(self, name) for name in pnames])
+
  @abstractmethod
  def forward(self):
  """Needs override."""
@@ -993,6 +1052,11 @@ class LayerNormLinear(TransformerEngineBaseModule):
  together with the output of the linear transformation.
  Example use case: residual connection for transformer module is
  taken post layernorm.
+ parameters_split : Tuple[str, ...], default = None
+ if a tuple of strings is provided, the weight and bias parameters of the
+ module are exposed as `N` separate `torch.nn.parameter.Parameter`s each,
+ split along the first dimension, where `N` is the length of the argument
+ and the strings contained are the names of the split parameters.
 
  Parallelism parameters
  ----------------------
@@ -1047,6 +1111,7 @@ def __init__(
  parallel_mode: Optional[str] = None,
  return_layernorm_output: bool = False,
  skip_weight_param_allocation: bool = False,
+ parameters_split: Optional[Tuple[str, ...]] = None,
  ) -> None:
  super().__init__()
  self.in_features = in_features
@@ -1055,7 +1120,7 @@ def __init__(
  self.use_bias = bias
  self.return_bias = return_bias
  self.return_layernorm_output = return_layernorm_output
- self.skip_weight_param_allocation = skip_weight_param_allocation
+ self.parameters_split = parameters_split
 
  if tp_group is None:
  self.tp_size = tp_size
@@ -1101,38 +1166,76 @@ def __init__(
  self.reset_layer_norm_parameters()
 
  if not skip_weight_param_allocation:
- self.weight = Parameter(
- torch.empty(
- self.out_features,
- self.in_features,
- device=torch.cuda.current_device(),
- dtype=params_dtype,
- )
- )
+ self.register_buffer("weight_tensor",
+ torch.empty(
+ self.out_features,
+ self.in_features,
+ device=torch.cuda.current_device(),
+ dtype=params_dtype),
+ persistent=False)
 
  initialize_affine_weight_gpu(
- self.weight,
+ self.weight_tensor,
  init_method,
  get_rng_state_tracker,
  partition_dim=1 if self.parallel_mode == "row" else 0,
  stride=1,
  )
 
  if self.use_bias or self.return_bias:
- self.bias = Parameter(
- torch.empty(
- self.out_features,
- device=torch.cuda.current_device(),
- dtype=params_dtype,
- )
- )
- if self.parallel_mode == "column":
- set_tensor_model_parallel_attributes(self.bias, True, 0, 1)
+ self.register_buffer("bias_tensor",
+ torch.empty(
+ self.out_features,
+ device=torch.cuda.current_device(),
+ dtype=params_dtype),
+ persistent=False)
  else:
- self.register_buffer("bias", torch.Tensor().type(params_dtype), persistent=False)
+ self.register_buffer(
+ "bias_tensor", torch.Tensor().type(params_dtype), persistent=False
+ )
 
  with torch.no_grad():
- self.bias.zero_()
+ self.bias_tensor.zero_()
+
+ if parameters_split is None:
+ parameters_split = ("",)
+
+ assert (
+ self.out_features % len(parameters_split) == 0
+ ), f"Weight and bias params cannot be split into {len(parameters_split)} parts"
+
+ split_size = self.out_features // len(parameters_split)
+
+ self.weight_names = []
+ self.bias_names = []
+
+ for i, pname in enumerate(parameters_split):
+ wname = pname + "weight"
+ bname = pname + "bias"
+
+ self.register_parameter(
+ wname, Parameter(self.weight_tensor[i * split_size : (i+1) * split_size])
+ )
+
+ set_tensor_model_parallel_attributes(
+ tensor=getattr(self, wname),
+ is_parallel=True,
+ dim=1 if parallel_mode == "row" else 0,
+ stride=1,
+ )
+
+ if self.use_bias or self.return_bias:
+ self.register_parameter(
+ bname, Parameter(self.bias_tensor[i * split_size : (i+1) * split_size])
+ )
+ else:
+ self.register_buffer(bname, torch.Tensor().type(params_dtype), persistent=False)
+
+ if parallel_mode == "column":
+ set_tensor_model_parallel_attributes(getattr(self, bname), True, 0, 1)
+
+ self.weight_names.append(wname)
+ self.bias_names.append(bname)
 
  self.fp8_weight_shapes.append(torch.Size((self.out_features, self.in_features)))
 
@@ -1193,7 +1296,18 @@ def forward(
  """
 
  with self.prepare_forward(inp, is_first_microbatch) as inp:
- bias_tensor = bias if bias is not None else self.bias
+ bias_tensor = (
+ bias if bias is not None
+ else self.bias if self.parameters_split is None
+ else self.bias_tensor if not self.training
+ else self.noop_cat("bias_tensor", self.bias_names)
+ )
+ weight_tensor = (
+ weight if weight is not None
+ else self.weight if self.parameters_split is None
+ else self.weight_tensor if not self.training
+ else self.noop_cat("weight_tensor", self.weight_names)
+ )
 
  if self.training:
  fwd_fn = _LayerNormLinear.apply
@@ -1205,7 +1319,7 @@ def forward(
  inp,
  self.layer_norm_weight,
  self.layer_norm_bias,
-weight if weight is not None else self.weight,
+weight_tensor,
  self.weight1_fp8 if self.fp8 else None,
  self.weight1_t_fp8 if self.fp8 else None,
  bias_tensor,
@@ -1586,6 +1700,11 @@ class Linear(TransformerEngineBaseModule):
  init_method : Callable, default = `None`
  used for initializing weights in the following way: `init_method(weight)`.
  When set to `None`, defaults to `torch.nn.init.normal_(mean=0.0, std=0.023)`.
+ parameters_split : Tuple[str, ...], default = None
+ if a tuple of strings is provided, the weight and bias parameters of the
+ module are exposed as `N` separate `torch.nn.parameter.Parameter`s each,
+ split along the first dimension, where `N` is the length of the argument
+ and the strings contained are the names of the split parameters.
 
  Parallelism parameters
  ----------------------
@@ -1641,14 +1760,15 @@ def __init__(
  params_dtype: torch.dtype = torch.float32,
  parallel_mode: Optional[str] = None,
  skip_weight_param_allocation: bool = False,
+ parameters_split: Optional[Tuple[str, ...]] = None,
  ) -> None:
  super().__init__()
  self.in_features = in_features
  self.out_features = out_features
  self.fuse_wgrad_accumulation = fuse_wgrad_accumulation
  self.use_bias = bias
  self.return_bias = return_bias
- self.skip_weight_param_allocation = skip_weight_param_allocation
+ self.parameters_split = parameters_split
 
  if tp_group is None:
  self.tp_size = tp_size
@@ -1675,38 +1795,76 @@ def __init__(
  self.sequence_parallel = (self.tp_size > 1) and sequence_parallel
 
  if not skip_weight_param_allocation:
- self.weight = Parameter(
- torch.empty(
- self.out_features,
- self.in_features,
- device=torch.cuda.current_device(),
- dtype=params_dtype,
- )
- )
+ self.register_buffer("weight_tensor",
+ torch.empty(
+ self.out_features,
+ self.in_features,
+ device=torch.cuda.current_device(),
+ dtype=params_dtype),
+ persistent=False)
 
  initialize_affine_weight_gpu(
- self.weight,
+ self.weight_tensor,
  init_method,
  get_rng_state_tracker,
  partition_dim=1 if self.parallel_mode == "row" else 0,
  stride=1,
  )
 
  if self.use_bias or self.return_bias:
- self.bias = Parameter(
- torch.empty(
- self.out_features,
- device=torch.cuda.current_device(),
- dtype=params_dtype,
- )
- )
- if self.parallel_mode == "column":
- set_tensor_model_parallel_attributes(self.bias, True, 0, 1)
+ self.register_buffer("bias_tensor",
+ torch.empty(
+ self.out_features,
+ device=torch.cuda.current_device(),
+ dtype=params_dtype),
+ persistent=False)
  else:
- self.register_buffer("bias", torch.Tensor().type(params_dtype), persistent=False)
+ self.register_buffer(
+ "bias_tensor", torch.Tensor().type(params_dtype), persistent=False
+ )
 
  with torch.no_grad():
- self.bias.zero_()
+ self.bias_tensor.zero_()
+
+ if parameters_split is None:
+ parameters_split = ("",)
+
+ assert (
+ self.out_features % len(parameters_split) == 0
+ ), f"Weight and bias params cannot be split into {len(parameters_split)} parts"
+
+ split_size = self.out_features // len(parameters_split)
+
+ self.weight_names = []
+ self.bias_names = []
+
+ for i, pname in enumerate(parameters_split):
+ wname = pname + "weight"
+ bname = pname + "bias"
+
+ self.register_parameter(
+ wname, Parameter(self.weight_tensor[i * split_size : (i+1) * split_size])
+ )
+
+ set_tensor_model_parallel_attributes(
+ tensor=getattr(self, wname),
+ is_parallel=True,
+ dim=1 if parallel_mode == "row" else 0,
+ stride=1,
+ )
+
+ if self.use_bias or self.return_bias:
+ self.register_parameter(
+ bname, Parameter(self.bias_tensor[i * split_size : (i+1) * split_size])
+ )
+ else:
+ self.register_buffer(bname, torch.Tensor().type(params_dtype), persistent=False)
+
+ if parallel_mode == "column":
+ set_tensor_model_parallel_attributes(getattr(self, bname), True, 0, 1)
+
+ self.weight_names.append(wname)
+ self.bias_names.append(bname)
 
  self.fp8_weight_shapes.append(torch.Size((self.out_features, self.in_features)))
 
@@ -1755,7 +1913,18 @@ def forward(
  """
 
  with self.prepare_forward(inp, is_first_microbatch) as inp:
- bias_tensor = bias if bias is not None else self.bias
+ bias_tensor = (
+ bias if bias is not None
+ else self.bias if self.parameters_split is None
+ else self.bias_tensor if not self.training
+ else self.noop_cat("bias_tensor", self.bias_names)
+ )
+ weight_tensor = (
+ weight if weight is not None
+ else self.weight if self.parameters_split is None
+ else self.weight_tensor if not self.training
+ else self.noop_cat("weight_tensor", self.weight_names)
+ )
 
  if self.training:
  linear_fn = _Linear.apply
@@ -1764,7 +1933,7 @@ def forward(
  linear_fn = _Linear.forward
  args = [None]
  args += (
-weight if weight is not None else self.weight,
+weight_tensor,
  self.weight1_fp8 if self.fp8 else None,
  self.weight1_t_fp8 if self.fp8 else None,
  inp,

transformer_engine/pytorch/module.py

@@ -7,7 +7,7 @@
 import pickle
 import warnings
 from abc import ABC, abstractmethod
-from typing import Union, Optional, Callable, Tuple, Dict, Any, Mapping
+from typing import Union, Optional, Callable, Tuple, Dict, Any, Mapping, List
 from functools import partial
 from contextlib import contextmanager
 
@@ -135,6 +135,42 @@ def _prepare_backward(fp8: bool,
  delete_key_from_amax_buffer(forward=False)
 
 
+class _NoopCat(torch.autograd.Function):
+ """This class is a no-op replacement for `torch.cat`."""
+
+ @staticmethod
+ def forward(ctx,
+ full_param_buffer: torch.Tensor,
+ *params_split: Tuple[torch.Tensor, ...],
+ ) -> torch.Tensor:
+ assert not full_param_buffer.requires_grad, "Buffers should not require gradient"
+ assert (
+ full_param_buffer.shape[0] % len(params_split) == 0
+ ), "Dimensions not compatible for concatenation"
+
+ param_temp = full_param_buffer.new()
+ param_temp.set_(full_param_buffer.storage(),
+ full_param_buffer.storage_offset(),
+ full_param_buffer.size(),
+ full_param_buffer.stride())
+ param_temp.requires_grad = True
+
+ ctx.save_for_backward(full_param_buffer, *params_split)
+ return param_temp
+
+ @staticmethod
+ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], ...]:
+ full_param_buffer, *params_split = ctx.saved_tensors
+
+ split_size = full_param_buffer.shape[0] // len(params_split)
+ grads = []
+
+ for i, _ in enumerate(params_split):
+ grads.append(grad_output[i * split_size : (i+1) * split_size])
+
+ return None, *grads
+
+
 class TransformerEngineBaseModule(torch.nn.Module, ABC):
  """Base TE module."""
 
@@ -572,6 +608,29 @@ def grad_output_preprocess(
 
  return grad_output_mat, grad_output_c, grad_output_t, grad_bias
 
+ def noop_cat(self, buffer_name: str, pnames: List[str]) -> torch.Tensor:
+ """No-op replacement of `torch.cat`. The buffer and split parameters must occupy
+ the same memory region. If this is not the case, then the split parameters
+ are concatenated and the buffer is overwritten. The parameters' memory is then
+ re-assigned to point to the buffer to avoid subsequent concatenations.
+ """
+
+ assert hasattr(self, buffer_name), f"No buffer named {buffer_name}"
+ full_param_buffer = getattr(self, buffer_name)
+ split_size = full_param_buffer.shape[0] // len(pnames)
+ params = [getattr(self, name) for name in pnames]
+ for i, p in enumerate(params):
+ if p.data.data_ptr() != full_param_buffer[i*split_size : (i+1)*split_size].data_ptr():
+ with torch.no_grad():
+ setattr(self, buffer_name, torch.cat(params))
+ for j, pname in enumerate(pnames):
+ full_param_buffer = getattr(self, buffer_name)
+ setattr(self, pname,
+ Parameter(full_param_buffer[j*split_size : (j+1)*split_size]))
+ break
+
+ return _NoopCat.apply(getattr(self, buffer_name), *[getattr(self, name) for name in pnames])
+
  @abstractmethod
  def forward(self):
  """Needs override."""
@@ -993,6 +1052,11 @@ class LayerNormLinear(TransformerEngineBaseModule):
  together with the output of the linear transformation.
  Example use case: residual connection for transformer module is
  taken post layernorm.
+ parameters_split : Tuple[str, ...], default = None
+ if a tuple of strings is provided, the weight and bias parameters of the
+ module are exposed as `N` separate `torch.nn.parameter.Parameter`s each,
+ split along the first dimension, where `N` is the length of the argument
+ and the strings contained are the names of the split parameters.
 
  Parallelism parameters
  ----------------------
@@ -1047,6 +1111,7 @@ def __init__(
  parallel_mode: Optional[str] = None,
  return_layernorm_output: bool = False,
  skip_weight_param_allocation: bool = False,
+ parameters_split: Optional[Tuple[str, ...]] = None,
  ) -> None:
  super().__init__()
  self.in_features = in_features
@@ -1055,7 +1120,7 @@ def __init__(
  self.use_bias = bias
  self.return_bias = return_bias
  self.return_layernorm_output = return_layernorm_output
- self.skip_weight_param_allocation = skip_weight_param_allocation
+ self.parameters_split = parameters_split
 
  if tp_group is None:
  self.tp_size = tp_size
@@ -1101,38 +1166,76 @@ def __init__(
  self.reset_layer_norm_parameters()
 
  if not skip_weight_param_allocation:
- self.weight = Parameter(
- torch.empty(
- self.out_features,
- self.in_features,
- device=torch.cuda.current_device(),
- dtype=params_dtype,
- )
- )
+ self.register_buffer("weight_tensor",
+ torch.empty(
+ self.out_features,
+ self.in_features,
+ device=torch.cuda.current_device(),
+ dtype=params_dtype),
+ persistent=False)
 
  initialize_affine_weight_gpu(
- self.weight,
+ self.weight_tensor,
  init_method,
  get_rng_state_tracker,
  partition_dim=1 if self.parallel_mode == "row" else 0,
  stride=1,
  )
 
  if self.use_bias or self.return_bias:
- self.bias = Parameter(
- torch.empty(
- self.out_features,
- device=torch.cuda.current_device(),
- dtype=params_dtype,
- )
- )
- if self.parallel_mode == "column":
- set_tensor_model_parallel_attributes(self.bias, True, 0, 1)
+ self.register_buffer("bias_tensor",
+ torch.empty(
+ self.out_features,
+ device=torch.cuda.current_device(),
+ dtype=params_dtype),
+ persistent=False)
  else:
- self.register_buffer("bias", torch.Tensor().type(params_dtype), persistent=False)
+ self.register_buffer(
+ "bias_tensor", torch.Tensor().type(params_dtype), persistent=False
+ )
 
  with torch.no_grad():
- self.bias.zero_()
+ self.bias_tensor.zero_()
+
+ if parameters_split is None:
+ parameters_split = ("",)
+
+ assert (
+ self.out_features % len(parameters_split) == 0
+ ), f"Weight and bias params cannot be split into {len(parameters_split)} parts"
+
+ split_size = self.out_features // len(parameters_split)
+
+ self.weight_names = []
+ self.bias_names = []
+
+ for i, pname in enumerate(parameters_split):
+ wname = pname + "weight"
+ bname = pname + "bias"
+
+ self.register_parameter(
+ wname, Parameter(self.weight_tensor[i * split_size : (i+1) * split_size])
+ )
+
+ set_tensor_model_parallel_attributes(
+ tensor=getattr(self, wname),
+ is_parallel=True,
+ dim=1 if parallel_mode == "row" else 0,
+ stride=1,
+ )
+
+ if self.use_bias or self.return_bias:
+ self.register_parameter(
+ bname, Parameter(self.bias_tensor[i * split_size : (i+1) * split_size])
+ )
+ else:
+ self.register_buffer(bname, torch.Tensor().type(params_dtype), persistent=False)
+
+ if parallel_mode == "column":
+ set_tensor_model_parallel_attributes(getattr(self, bname), True, 0, 1)
+
+ self.weight_names.append(wname)
+ self.bias_names.append(bname)
 
  self.fp8_weight_shapes.append(torch.Size((self.out_features, self.in_features)))
 
@@ -1193,7 +1296,18 @@ def forward(
  """
 
  with self.prepare_forward(inp, is_first_microbatch) as inp:
- bias_tensor = bias if bias is not None else self.bias
+ bias_tensor = (
+ bias if bias is not None
+ else self.bias if self.parameters_split is None
+ else self.bias_tensor if not self.training
+ else self.noop_cat("bias_tensor", self.bias_names)
+ )
+ weight_tensor = (
+ weight if weight is not None
+ else self.weight if self.parameters_split is None
+ else self.weight_tensor if not self.training
+ else self.noop_cat("weight_tensor", self.weight_names)
+ )
 
  if self.training:
  fwd_fn = _LayerNormLinear.apply
@@ -1205,7 +1319,7 @@ def forward(
  inp,
  self.layer_norm_weight,
  self.layer_norm_bias,
-weight if weight is not None else self.weight,
+weight_tensor,
  self.weight1_fp8 if self.fp8 else None,
  self.weight1_t_fp8 if self.fp8 else None,
  bias_tensor,
@@ -1586,6 +1700,11 @@ class Linear(TransformerEngineBaseModule):
  init_method : Callable, default = `None`
  used for initializing weights in the following way: `init_method(weight)`.
  When set to `None`, defaults to `torch.nn.init.normal_(mean=0.0, std=0.023)`.
+ parameters_split : Tuple[str, ...], default = None
+ if a tuple of strings is provided, the weight and bias parameters of the
+ module are exposed as `N` separate `torch.nn.parameter.Parameter`s each,
+ split along the first dimension, where `N` is the length of the argument
+ and the strings contained are the names of the split parameters.
 
  Parallelism parameters
  ----------------------
@@ -1641,14 +1760,15 @@ def __init__(
  params_dtype: torch.dtype = torch.float32,
  parallel_mode: Optional[str] = None,
  skip_weight_param_allocation: bool = False,
+ parameters_split: Optional[Tuple[str, ...]] = None,
  ) -> None:
  super().__init__()
  self.in_features = in_features
  self.out_features = out_features
  self.fuse_wgrad_accumulation = fuse_wgrad_accumulation
  self.use_bias = bias
  self.return_bias = return_bias
- self.skip_weight_param_allocation = skip_weight_param_allocation
+ self.parameters_split = parameters_split
 
  if tp_group is None:
  self.tp_size = tp_size
@@ -1675,38 +1795,76 @@ def __init__(
  self.sequence_parallel = (self.tp_size > 1) and sequence_parallel
 
  if not skip_weight_param_allocation:
- self.weight = Parameter(
- torch.empty(
- self.out_features,
- self.in_features,
- device=torch.cuda.current_device(),
- dtype=params_dtype,
- )
- )
+ self.register_buffer("weight_tensor",
+ torch.empty(
+ self.out_features,
+ self.in_features,
+ device=torch.cuda.current_device(),
+ dtype=params_dtype),
+ persistent=False)
 
  initialize_affine_weight_gpu(
- self.weight,
+ self.weight_tensor,
  init_method,
  get_rng_state_tracker,
  partition_dim=1 if self.parallel_mode == "row" else 0,
  stride=1,
  )
 
  if self.use_bias or self.return_bias:
- self.bias = Parameter(
- torch.empty(
- self.out_features,
- device=torch.cuda.current_device(),
- dtype=params_dtype,
- )
- )
- if self.parallel_mode == "column":
- set_tensor_model_parallel_attributes(self.bias, True, 0, 1)
+ self.register_buffer("bias_tensor",
+ torch.empty(
+ self.out_features,
+ device=torch.cuda.current_device(),
+ dtype=params_dtype),
+ persistent=False)
  else:
- self.register_buffer("bias", torch.Tensor().type(params_dtype), persistent=False)
+ self.register_buffer(
+ "bias_tensor", torch.Tensor().type(params_dtype), persistent=False
+ )
 
  with torch.no_grad():
- self.bias.zero_()
+ self.bias_tensor.zero_()
+
+ if parameters_split is None:
+ parameters_split = ("",)
+
+ assert (
+ self.out_features % len(parameters_split) == 0
+ ), f"Weight and bias params cannot be split into {len(parameters_split)} parts"
+
+ split_size = self.out_features // len(parameters_split)
+
+ self.weight_names = []
+ self.bias_names = []
+
+ for i, pname in enumerate(parameters_split):
+ wname = pname + "weight"
+ bname = pname + "bias"
+
+ self.register_parameter(
+ wname, Parameter(self.weight_tensor[i * split_size : (i+1) * split_size])
+ )
+
+ set_tensor_model_parallel_attributes(
+ tensor=getattr(self, wname),
+ is_parallel=True,
+ dim=1 if parallel_mode == "row" else 0,
+ stride=1,
+ )
+
+ if self.use_bias or self.return_bias:
+ self.register_parameter(
+ bname, Parameter(self.bias_tensor[i * split_size : (i+1) * split_size])
+ )
+ else:
+ self.register_buffer(bname, torch.Tensor().type(params_dtype), persistent=False)
+
+ if parallel_mode == "column":
+ set_tensor_model_parallel_attributes(getattr(self, bname), True, 0, 1)
+
+ self.weight_names.append(wname)
+ self.bias_names.append(bname)
 
  self.fp8_weight_shapes.append(torch.Size((self.out_features, self.in_features)))
 
@@ -1755,7 +1913,18 @@ def forward(
  """
 
  with self.prepare_forward(inp, is_first_microbatch) as inp:
- bias_tensor = bias if bias is not None else self.bias
+ bias_tensor = (
+ bias if bias is not None
+ else self.bias if self.parameters_split is None
+ else self.bias_tensor if not self.training
+ else self.noop_cat("bias_tensor", self.bias_names)
+ )
+ weight_tensor = (
+ weight if weight is not None
+ else self.weight if self.parameters_split is None
+ else self.weight_tensor if not self.training
+ else self.noop_cat("weight_tensor", self.weight_names)
+ )
 
  if self.training:
  linear_fn = _Linear.apply
@@ -1764,7 +1933,7 @@ def forward(
  linear_fn = _Linear.forward
  args = [None]
  args += (
-weight if weight is not None else self.weight,
+weight_tensor,
  self.weight1_fp8 if self.fp8 else None,
  self.weight1_t_fp8 if self.fp8 else None,
  inp,