如何处理序列模型中极高规模的目标词汇量问题？

我之前在做超大词汇量的序列模型时，也碰到过一模一样的GPU内存爆炸问题——训练用sampled softmax没问题，但评估时全量softmax直接把卡干爆了。给你分享几个亲测有效的解决方案，按优先级排序：

直接计算80M维度的softmax对GPU来说太苛刻了，我们可以把词汇表拆成若干小批次，分块计算logits，最后再拼接结果。这样每次GPU只处理一小部分词汇的权重，内存压力直接降下来。

举个TensorFlow的代码示例：

vocab_size = 80000000
split_size = 100000  # 每次处理10万词汇，可根据GPU内存调整
num_splits = (vocab_size + split_size - 1) // split_size

# 假设你的模型最后一层是Dense层，权重可以通过model.layers[-1].weights获取
kernel, bias = model.layers[-1].weights

# 分批计算logits（hidden_states是模型的最后一层隐藏态输出）
batch_logits = []
for i in range(num_splits):
    start_idx = i * split_size
    end_idx = min((i+1)*split_size, vocab_size)
    # 取当前批次的权重
    batch_kernel = kernel[:, start_idx:end_idx]
    batch_bias = bias[start_idx:end_idx]
    # 计算当前批次的logits
    current_logits = tf.matmul(hidden_states, batch_kernel) + batch_bias
    batch_logits.append(current_logits)

# 拼接所有批次的logits得到全量logits
full_logits = tf.concat(batch_logits, axis=-1)
# 计算全量softmax（如果确实需要的话）
full_softmax = tf.nn.softmax(full_logits)

如果只是计算perplexity，其实不需要全量softmax——你只需要真实标签对应的log概率，这时候可以在分批计算时直接提取对应位置的logits，再结合分批计算的logsumexp来得到最终的log概率，进一步节省内存。

如果GPU内存实在捉襟见肘，评估阶段直接把计算转到CPU上就行。服务器的CPU内存通常比GPU大得多，能轻松容纳80M维度的张量，唯一的缺点是速度慢一点，但评估阶段不像训练那样频繁，完全可以接受。

代码示例：

# 评估时指定CPU设备
with tf.device('/CPU:0'):
    full_logits = model(hidden_states, training=False)
    full_softmax = tf.nn.softmax(full_logits)
    # 后续计算评估指标...

很多时候我们误以为必须计算全量softmax，但实际上大部分评估指标可以绕开这一步：

• Perplexity计算：只需要每个样本真实标签的log概率，而log概率 = logits[label_idx] - logsumexp(full_logits)。其中logsumexp可以分批计算：先对每个词汇块的logits计算logsumexp，再对这些中间结果计算一次logsumexp，这样就不用一次性加载全量logits到GPU。
• Top-K准确率：可以分批计算每个词汇块的logits，维护全局的top-k值，最后统计命中情况，同样不需要全量softmax。

比如分批计算logsumexp的代码：

def batch_logsumexp(logits, split_size=100000):
    num_splits = (logits.shape[-1] + split_size - 1) // split_size
    split_logits = tf.split(logits, num_splits, axis=-1)
    block_lse = [tf.reduce_logsumexp(split, axis=-1) for split in split_logits]
    # 合并各块的logsumexp
    return tf.reduce_logsumexp(tf.stack(block_lse, axis=-1), axis=-1)

# 计算真实标签的log概率
label_logits = tf.gather(full_logits, labels, axis=-1, batch_dims=1)
lse = batch_logsumexp(full_logits)
log_probs = label_logits - lse
# 用log_probs计算perplexity
perplexity = tf.exp(tf.reduce_mean(-log_probs))

开启TensorFlow的混合精度，将大部分张量从float32转为float16，内存占用直接减半，这对大张量的场景非常友好。需要注意的是，logsumexp这类容易出现数值溢出的操作，最好保留float32计算，避免精度损失。

代码示例：

import tensorflow as tf
tf.keras.mixed_precision.set_global_policy('mixed_float16')

# 构建你的模型...
# 注意：最后一层Dense的输出可以用float32，避免softmax的数值问题
model.add(tf.keras.layers.Dense(vocab_size, dtype='float32'))

如果有多个GPU，可以把最后一层Dense的权重分片到不同GPU上，每个GPU负责一部分词汇的logits计算，最后合并结果。TensorFlow的tf.distribute.MirroredStrategy可以自动处理部分并行逻辑，或者你也可以手动分片权重，在不同GPU上计算后拼接。

内容的提问来源于stack exchange，提问作者user_12