Huggingface Transformers实现张量并行的小坑 set/get_output_embeddings

transformers 库里实现的很多模型会有这么两个函数 get_output_embeddings和 get_output_embeddings。以 SwitchTransformer 为例

class SwitchTransformersForConditionalGeneration(SwitchTransformersPreTrainedModel):
    def set_output_embeddings(self, new_embeddings):
        self.lm_head = new_embeddings
    def get_output_embeddings(self):
        return self.lm_head

默认情况下，大模型的输入和输出的 vocab 是保持一致的，所以如果传入的 embedding 的大小变化了，默认也会让 lm_head 发生变化。

但是在实现张量并行的时候，我们通常会使用如下方式来初始化lm_head

from fairscale.nn.model_parallel.layers import (
    ParallelEmbedding,
    RowParallelLinear,
    ColumnParallelLinear
)
default_linear_init = functools.partial(nn.init.kaiming_uniform_, a=math.sqrt(5))
def __init__(self, ...):
    self.lm_head = ColumnParallelLinear(config.d_model, config.vocab_size, bias=False, init_method=default_linear_init)

换言之，在多 GPU 张量并行下，每张卡上 lm_head 的输出维度就不再是原来的 vocab_size 了，而是 vocab_size/#gpus。所以一种粗暴的解决办法就是把get_output_embeddings的输出改为 None 即可，如下：

    def get_output_embeddings(self):
        return None # PretrainedModel.tie_weights 函数会将 lm_head 绑定为 shared 参数，导致张量并行情况下 lm_head 参数发生不匹配的错误