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 參數發生不匹配的錯誤