TiedLinear

class torchtune.modules.TiedLinear(tied_module: Module)

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 qwen2_0_5b(), qwen2_1_5b() and all of the gemma() models. It requires as input an nn.Module, instead of the weight of the module, so it can work with FSDP. Otherwise, the memory reference will be lost after FSDP is applied.

Parameters:

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’.