Weight tying between embedding and LM head layer

[Luosuu]

Hi TorchTitan team,

I am wondering can simpleFSDP support weight tying which is also a thing in Llama-3? I did some initial exploration but it seems not supported yet?

Thank you very much!

[tianyu-l]

Currently there's no plan to support weight tying. Since the source code is in the repo, you are encouraged to try it by yourself.
First step may be: don't repeatedly do sharding for both embedding weight and output weight
https://github.com/pytorch/torchtitan/blob/main/torchtitan/experiments/simple_fsdp/simple_fsdp.py#L295

[HaniItani]

Hi @tianyu-l ,

I'm also interested in weight tying. If I naively remove the sharding of the output layer from infra/parallelize.py, I encounter RuntimeError: aten.mm.default: got mixed torch.Tensor and DTensor, need to convert all torch.Tensor to DTensor before calling distributed operators!.

I'm trying to follow the torchtune implementation of weight tying. They do this:

class Linear(nn.Module):
    """
    nn.Module used in :func:`~torchtune.modules.tied_linear.TiedLinear`, added to work with the hooks
    :class:`~torchtune.training._activation_offloading.NoOpManager` that ignore activation
    offloading context manager.

    Without this class, we can't add NoOp hooks, and we will offload the activation of
    the tied linear layer, which is slow.

    For more information, see how NoOpManager is called in the recipes.
    """

    def forward(self, x: torch.Tensor, weight: torch.Tensor):
        return F.linear(x, weight)


class TiedLinear:
    """
    A tied linear layer, without bias, that shares the same weight as another linear layer.
    This is useful for models that use tied weights, such as :func:`~torchtune.models.qwen2_0_5b`,
    :func:`~torchtune.models.qwen2_1_5b` and all of the :func:`~torchtune.models.gemma` and
    :func:`~torchtune.models.llama3_2` models.

    It requires as input an nn.Module, instead of the weight of the module, so it
    can work with FSDP. When FSDP is applied, the memory pointer to the weight is different,
    but the nn.Module remains the same. This is why we need to pass the nn.Module instead of
    the weight, if we want to keep the weights tied.

    Args:
        tied_module (nn.Module): The module whose weight is shared. Only
            the weight is used. The bias is ignored.
    Raises:
        AttributeError: If the provided module does not have an attribute 'weight'.
    """

    def __init__(self, tied_module: nn.Module):
        self.tied_module = tied_module
        self.linear = Linear()
        if not hasattr(tied_module, "weight"):
            raise AttributeError(
                "Provided module does not have attribute 'weight'. Please check your tied_module."
            )

    def __call__(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x (torch.Tensor): Input tensor. Should have shape ``(..., in_dim)``, where ``in_dim``
                is the input dimension of the tied module.
        Returns:
            torch.Tensor: The output tensor, having shape ``(..., out_dim)``, where ``out_dim`` is \
                the output dimension of the tied module.
        """
        return self.linear(x, self.tied_module.weight)

And remove sharding of embedding/output from infra/parallelize.py. Do you think this is at least on the right track or I'm missing something? The code runs, but the starting loss is around 1000 vs 12 (without tying word embeddings)

Would really appreciate your input and guidance. Thanks!