Megatron-LM中的权重初始化
现象描述:在某平台上上训练大模型时,受限于内存不够,需要使用模型并行。
例如模型结构:下采样(包含卷积等层)⇒SwinTransformer⇒ 上采样(包含卷积等),Megatron-LM只对Transformer部分进行了初始化。
当tp=6时,如果全部随机初始化权重,transformer层在不同rank上的权重是不同的(因为每个rank保持不同部分transformer权重);但其他无法并行的层(例如卷积) (经打印验证) 在6个节点都保持了相同的初始化权重。
由此对Megatron-LM和PyTorch中权重的初始化方式产生了疑惑。
随机初始化
我们以torch.nn.Conv2d为例子,它的随机初始化定义在基类
torch.nn.modules.conv._ConvNd
torch/nn/modules/conv.py
- 可以Google搜索:caffe2.ai: xxxtorch api 来查找
Caffe2 - Python API: torch.nn.modules.conv._ConvNd Class Reference
nn.Module是所有神经网络单元(neural network modules)的基类,其中负责初始化参数的函数是 reset_parameters(self)
class _ConvNd(Module):
__constants__ = ['stride', 'padding', 'dilation', 'groups', 'bias', 'padding_mode']
def __init__(self, in_channels, out_channels, kernel_size, stride,
padding, dilation, transposed, output_padding,
groups, bias, padding_mode):
......
if transposed:
self.weight = Parameter(torch.Tensor(
in_channels, out_channels // groups, *kernel_size))
else:
self.weight = Parameter(torch.Tensor(
out_channels, in_channels // groups, *kernel_size))
if bias:
self.bias = Parameter(torch.Tensor(out_channels))
else:
self.register_parameter('bias', None)
self.reset_parameters()
......
def reset_parameters(self):
n = self.in_channels
init.kaiming_uniform_(self.weight, a=math.sqrt(5))
if self.bias is not None:
fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
bound = 1 / math.sqrt(fan_in)
init.uniform_(self.bias, -bound, bound)
可以看到,随机参数的初始化分为如下两个步骤:
self.weight = Parameter( torch.Tensor(out_channels, in_channels // group, **kernel_size ) )
self.bias = Parameter(torch.Tensor(out_channels))
这里torch.nn.Parameter 继承自 torch.Tensor,使用Parameter创建的张量:
- 会自动设置 requires_grad 为True。
- 会将该变量加入到 模型的named_parameters当中。
当然,如果没有
reset_parameters 对 self.weight 和 self.bias进行重制。
- weight初始化: init.kaiming_uniform_(self.weight, a=math.sqrt(5)),这里使用的是HE初始化(凯明何的初始化)的均匀分布采样,采样的bound:
$$
bound=\sqrt{\frac{6}{(1+a)^2\times fan_{in}}}
$$其中$fan_{in}$ 是输入神经元数量,输入通道 x 卷积大小K^2
- bias初始化:首先根据weight来计算$fan_{in}$和$fan_{out}$,然后用
$$
bound=\sqrt{\frac{1}{fan_{in}}}
$$采用均匀分布进行采样。
这里,bias和weight都在相同的采样区间里,采用均匀分布来采样。
关于torch.nn.init中的采样方法,可以参见blog:
那么我们来关注另一个问题,有了初始化方法,它又是怎么同随机种子相关联的?
首先,megatron的随机种子是在args.seed参数来设置的,默认1234。
但实际设置时,采用了函数_set_random_seed
def _set_random_seed(seed_, data_parallel_random_init=False):
"""Set random seed for reproducability."""
if seed_ is not None and seed_ > 0:
# Ensure that different pipeline MP stages get different seeds.
seed = seed_ + (100 * mpu.get_pipeline_model_parallel_rank())
# Ensure different data parallel ranks get different seeds
if data_parallel_random_init:
seed = seed + (10 * mpu.get_data_parallel_rank())
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
if torch.cuda.device_count() > 0:
tensor_parallel.model_parallel_cuda_manual_seed(seed)
else:
raise ValueError('Seed ({}) should be a positive integer.'.format(seed))
这里有些细节:
- 流水线并行组各rank上的seed是不同的,那么流水线并行不同rank上 哪怕是未切分的层也是不同的随机参数。 但是这保证了 不同Transformer层的权重是不一样的。
- 首先,默认情况下,不同的流水线阶段 之间的 种子是不一样的,容易理解 不同transformer层的初始化种子不一致。
- 可选的,可以让不同数据对应的模型副本采用不同的随机初始化方法。
- 设置了三种种子:
- random.seed:影响random.random… random.randint等
- np.random.seed:影响np.random.rand()、np.random.randn()等
- torch.manual_seed:影响torch.rand()、torch.randn()和torch.randint()等。
这意味着:
张量并行组的各rank上的seed是相同的,也就意味着 随机发生器所生成的随机序列是确定的,对于非切分的层,对应层的初始化参数是完全相同的。
- 但是不同次序组织的层,哪怕相同规模和设置,由于随机生成的发生次序不一样,因此彼此之间随机参数不同。
- 而不同规模或者类型的层,哪怕两次运行位于相同的随机次序上,由于采样的范围、采样的方法不同,生成的随机参数也是不一样的。
- 实验证明,torch.init..随机方法,每次会消耗一个随机生成器产生的随机状态,比如在一个卷积层中,初始化了两个参数,消耗了两个随机状态, 下一个卷积或者层的初始化参数会使用第三个状态。 那么,我们可以用两个torch.init.去替代第一个卷积,保证下一个卷积生成的随机参数还是相同的。
import torch import random import numpy as np seed=1234 random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) #base和1_1 都能保证bias生成相同, layer2weight相同· #layer1 = torch.nn.Conv2d(10, 10, kernel_size=4) #bias shape: c_out 10, weight shape (c_out, c_in, k1, k2) = 1600, base #layer1 = torch.nn.Conv2d(40, 10, kernel_size=2) #1_1 bias shape: c_out 10, weight shape (10,40, 2,2) = 1600,相同bias采样范围, 相同的采样数量 layer1 = torch.nn.Conv2d(20, 10, kernel_size=2) #1_2 对照组,不同的采样范围和采样数量,但是layer2的权重相同 #猜测,一次完整采样用一个随机生成的东东 # 那加了layer3 和 两次随机采样应该是等效的,看layer2的weight是相同的 #layer3 = torch.nn.Conv2d(20, 10, kernel_size=2) #weight = torch.ones([10,20,2,2]) #bias = torch.ones([1,10]) #torch.nn.init.kaiming_uniform_(weight) #torch.nn.init.uniform_(bias) # ##进一步测试 增加一行bias的uniform,layer2的权重应该会有变化 #bias2 = torch.ones([1,10]) #torch.nn.init.uniform_(bias2) ##再进一步测试,使用两次init_uniform_,让layer2保持相同 #bias2 = torch.ones([1,10]) #torch.nn.init.uniform_(bias2) #torch.nn.init.uniform_(bias2) layer2 = torch.nn.Conv2d(3,3, kernel_size=2) #same_weight = torch.equal(layer1.weight,layer2.weight) #same_bias = torch.equal(layer1.bias, layer2.bias) #print(f"weight same is {same_weight}, bias same is {same_bias}") print(f"layer1.bias {layer1.bias}, layer2.weight {layer2.weight}")
(我们可以尝试,在随机生成参数时,让并行卷积层在相同的随机生成次序上,在相同采样范围和采样方法? )
那么对于张量并行,位于不同rank上的切分的层,是如何实现权重不同的?
我们可以在megatron-sw/megatron/core/tensor_parallel/layers.py
当中ColumnParallelLinear或者 RowParallelLinear找到答案:
if use_cpu_initialization:
self.weight = Parameter(torch.empty(self.output_size,
self.input_size_per_partition,
dtype=params_dtype))
if perform_initialization:
self.master_weight = _initialize_affine_weight_cpu(
self.weight, self.output_size, self.input_size,
self.input_size_per_partition, 1, init_method,
stride=stride, return_master_weight=keep_master_weight_for_test,
params_dtype=params_dtype)
对于weight而言,生成的是一个列被切分后大小,cpu初始化的话,这里还有一个master_weight
def _initialize_affine_weight_cpu(weight, output_size, input_size,
per_partition_size, partition_dim,
init_method, stride=1,
return_master_weight=False,
*, params_dtype=torch.float32):
"""Initialize affine weight for model parallel.
Build the master weight on all processes and scatter
the relevant chunk."""
set_tensor_model_parallel_attributes(tensor=weight,
is_parallel=True,
dim=partition_dim,
stride=stride)
# Initialize master weight
master_weight = torch.empty(output_size, input_size,
dtype=torch.float,
requires_grad=False)
init_method(master_weight)
master_weight = master_weight.to(dtype=params_dtype)
# Split and copy
per_partition_per_stride_size = divide(per_partition_size, stride)
weight_list = torch.split(master_weight, per_partition_per_stride_size,
dim=partition_dim)
rank = get_tensor_model_parallel_rank()
world_size = get_tensor_model_parallel_world_size()
my_weight_list = weight_list[rank::world_size]
with torch.no_grad():
torch.cat(my_weight_list, dim=partition_dim, out=weight)
if return_master_weight:
return master_weight
return None
这里可以看到,创建了一个master_weight权重副本,然后对它进行了切分,根据rank来返回不同的weight_list切片weight_list[rank::world_size]
,这里的start=rank,step=world_size. 貌似是比较安全的写法?但如果 切分份数就是world_size的情况下,显然是取 了rank对应的权重切片。
然后再将本rank的切片weight重新用cat的方式组织,覆盖写回weight当中。
但是在gpu版本的初始化中,是直接对weight进行了初始化。
def _initialize_affine_weight_gpu(weight, init_method,
partition_dim, stride=1):
"""Initialize affine weight for model parallel on GPU."""
set_tensor_model_parallel_attributes(tensor=weight,
is_parallel=True,
dim=partition_dim,
stride=stride)
with get_cuda_rng_tracker().fork():
init_method(weight)
这里保证了每个gpu(哪怕在随机种子相同的情况下)单独创建一个互不相同的独立RNG状态,保证彼此之间生成的是不一样的。
总结下,张量并行初始化线性层参数的过程:
GPU上,通过fork一个rng状态, 让不同GPU分别产生不同的随机参数;
CPU上,通过大家同时生成一个大的相同的随机参数副本,然后根据张量并行rank来选择不同的切片作为自己的随机参数。
张量并行中,(非切分)层的随机初始化过程是完全一致的,因为不同节点的种子一样,随机初始化的次序也相同。
但是流水线并行,会在megatron设置随机种子时,让不同流水线阶段rank的随机种子不同。