在Galvatron中添加新模型
本指南将教你如何在Galvatron中添加新模型。
目录结构
一个模型在Galvatron中的目录结构如下；
MyModel/
├── meta_configs/ # 模型配置文件目录
│ ├── __init__.py
│ ├── config_utils.py # 配置工具函数
│ ├── MyModel-{MODEL_SIZE}b.json # 模型配置
│ └── ... # 其他规模模型配置
│
├── scripts/ # 运行脚本目录
│ ├── profile.sh # 性能分析脚本
│ ├── train.sh # 训练脚本
│ └── search.sh # 并行策略搜索脚本
│
├── __init__.py
├── arguments.py # 参数定义
├── dataloader.py # 数据加载实现
├── profiler.py # 性能分析入口
├── search_dist.py # 并行策略搜索入口
├── train.py # 单机训练入口
├── train_dist.py # 分布式训练入口
├── train_dist_random.py # 随机数据训练入口
│
├── MyModelModel_checkpoint.py # 检查点保存加载
├── MyModelModel_hybrid_parallel.py # 混合并行实现
├── MyModelModel_sequential.py # 序列化模型实现
└── MyModelModel_tensor_parallel.py # 张量并行实现
Galvatron构建混合并行模型过程
在介绍如何加入新模型之前，我们先来了解一下Galvatron构建混合并行模型的大致过程。
Galvatron构建模型不需要手动定义模型整体结构，而是通过使用transformers或flash attention中相应的模型结构，你可以在MyModel中添加hf或fa后缀来区分你所选择的模型结构后端。如果你不知道该选择什么样的模型结构后端，我们推荐你选择hf，因为Galvatron对hf的支持更加全面（fa模型不支持Ulysses-SP并行方法）。接着基于得到的模型结构构件混合并行模型的流程在construct_hybrid_parallel_model_api中。其具体的流程如下：
预处理配置：获取混合并行策略、模型配置等信息
通信组生成 （Step 0）：生成各种并行策略需要的通信组
构建张量并行模型 （Step 1）：使用模型特定的 TP 函数（定义在
MyModelModel_tensor_parallel.py中）构建张量并行模型构建序列模型 （Step 2）：使用模型特定的序列化函数重构模型（定义在
MyModelModel_sequential.py中）包装重分布模块 （Step 3）：为模型添加数据重分布功能，保证每层的数据分布和并行策略对应
构建流水线并行 （Step 4）：构建流水线并行模型，将不同的stage放置在对应设备上
包装数据并行模块 （Step 5）：基于FSDP库包装数据并行模块
添加检查点包装 （Step 6）：根据检查点配置为模块添加检查点功能
其中，只有该API的调用，以及Step1和Step2实现需要使用模型特定的函数完成，其他步骤都是Galvatron的通用实现。
核心文件说明
添加新模型的核心是模型实现文件，这是开发者需要实现的最主要的部分，它定义了模型的结构和实现。
1 张量并行实现
张量并行实现通过MyModelModel_tensor_parallel.py文件实现，该文件定义了模型的张量并行实现，需要将Sequential中的模块替换成支持张量并行的模块，这里Galvatron根据不同的模型后端，提供了不同的张量并行实现，具体来说，hf使用Megatron-TP，fa使用flash-attn提供的TP。
对于hf，你需要实现MyModelLayer_tp类，并实现MyModelAttention_tp和MyModelMLP_tp类，对于fa，则可以直接调用flash_attn的create_mixer_cls和 create_mlp_cls方法。同时你还需要定义construct_tensor_parallel_model函数，用于将完整模型进行TP模型替换。这方面的详细例子可以参考gpt_hf和gpt_fa。
1.1 Transformer层 （hf模型格式）
Transformer层通过MyModelLayer_tp类实现:
class MyModelLayer_tp(nn.Module):
def __init__(self, config, layer_number, tp_group=None, sp_group=None):
"""
参数:
config: 模型配置对象，TransformerConfig
layer_number: 当前层的索引编号
tp_group: 当前层张量并行通信组，CommGroup
sp_group: 当前层序列并行通信组，CommGroup
"""
super().__init__()
self.attention = MyModelAttention_tp(config, layer_number, tp_group, sp_group)
self.mlp = MyModelMLP_tp(config, tp_group)
self.idx = layer_number
def forward(self, hidden_states, attention_mask=None):
# ...
pass
该类主要负责定义一层Transformer的实现，包括注意力机制和前馈神经网络，需要注意的是self.idx的定义是必要的，这关乎后面如何区分层，config则直接使用创建Transformer库中的模型时使用的TransformerConfig类。
1.2 注意力层 （hf模型格式）
注意力层通过MyModelAttention_tp类实现:
class MyModelAttention_tp(nn.Module):
def __init__(self, config, layer_number, tp_group=None, sp_group=None):
"""
参数:
config: 模型配置对象，TransformerConfig
layer_number: 当前层的索引编号
tp_group: 张量并行通信组，CommGroup
sp_group: 序列并行通信组，CommGroup
"""
super().__init__()
# ...
megatron_config = core_transformer_config_from_args(args)
self.attention = ParallelAttention(megatron_config, ...)
# ...
def forward(self, hidden_states, attention_mask):
# ...
pass
ParallelAttention是Galvatron修改后的Megatron-TP中的注意力层实现，在原版Megatron-TP的注意力层实现中，增加了tp_group、sp_group、use_ulysses三个参数，分别表示张量并行通信组、序列并行通信组、是否使用Ulysses序列并行，通常来说你可以直接参考gpt_hf的例子实现这部分。
1.3 前馈神经网络层（hf模型格式）
前馈神经网络层通过MyModelMLP_tp类实现:
class MyModelMLP_tp(nn.Module):
def __init__(self, config, tp_group=None):
"""
参数:
config: 模型配置对象，TransformerConfig
tp_group: 张量并行通信组，CommGroup
"""
super().__init__()
# ...
megatron_config = core_transformer_config_from_args(get_args())
self.mlp = ParallelMLP(megatron_config, tp_group = self.tp_group)
# ...
def forward(self, hidden_states):
# ...
pass
ParallelMLP是Galvatron修改后的Megatron-TP中的前馈神经网络层实现，在原版Megatron-TP的注意力层实现中，增加了tp_group这个参数，用于表示张量并行通信组，通常来说你可以直接参考gpt_hf的例子实现这部分。
1.4 构造张量并行模型（hf模型格式）
构造张量并行模型通过construct_tensor_parallel_model函数实现:
def construct_tensor_parallel_model(model, config, tp_groups_enc, sp_groups_enc):
"""
将模型转换为张量并行版本
参数:
model: 原始模型实例
config: 模型配置对象，TransformerConfig
tp_groups_enc: 每一层的张量并行通信组列表，List[CommGroup]
sp_groups_enc: 每一层的序列并行通信组列表，List[CommGroup]
返回:
转换后的张量并行模型
"""
# ...
pass
该函数主要完成三件事：将模型中的Transformer Layer替换为MyModelLayer_tp，将模型中的embedding层替换为VocabParallelEmbedding，将模型中的lm_head替换为ColumnParallelLinear。VocabParallelEmbedding和ColumnParallelLinear是同样是Galvatron修改后的Megatron-TP中的嵌入层和线性层实现，增加了tp_group和sp_group这两个参数，用于表示张量并行通信组和序列并行通信组，你也可以直接参考gpt_hf的例子实现这部分。
注意：这些类和函数中用到的通信组是Galvatron自定义的CommGroup类，如果你想访问torch生成的通信组，请使用tp_group.group和sp_group.group。
1.5 构造张量并行模型（fa模型格式）
对于fa，你只需要实现construct_tensor_parallel_model函数即可，在该函数中你需要将Transformer Layer中的attention和mlp模块分别替换为flash_attn的create_mixer_cls和 create_mlp_cls方法，将embedding层替换为flash_attn的ParallelGPT2Embeddings方法，将lm_head替换为flash_attn的ColumnParallelLinear方法。详细的例子请参考gpt_fa。
2 序列化模型实现
MyModelModel_sequential.py定义了模型的序列化实现，包括模型的前向传播和反向传播实现。
对于传统的Transformer模型，你需要实现MyModelEmbeddings_, MyModelLayers_, MyModelPreNorm_, MyModelCls_ 等类。
此外，还需要实现construct_sequential_model函数，用于将模型转换为序列化模型。以及MyModelModelInfo类，用于定义模型相关信息。
具体来说，每个类的定义和格式如下：
2.1 嵌入层
嵌入层通过MyModelEmbeddings_类实现:
class MyModelEmbeddings_(nn.Module):
def __init__(self, model):
"""
参数:
model: 模型实例
"""
super().__init__()
# ...
def forward(self, tokens, **kwargs):
# ...
pass
该类主要用于定义模型中的嵌入层，包括词嵌入、位置嵌入等。
这里__init__函数中需要传入的model是直接通过调用transformers或flash-attn获取到的模型（所有API中model都需要传入transformers或flash-attn获取到的模型）。
为了增强代码的健壮性，该函数还需要支持一些额外的特性：Megatron序列并行、Ulysses序列并行（fa不支持）,这方面的详细例子可以参考gpt_hf和gpt_fa。
注意：当使用hf后端时，对于有多种Embedding类型的文件（比如GPT同时拥有Vocab和Position Embedding），需要额外定义不同的Embedding类以区分这两种不同的Embedding参数，gpt_hf中展示了这样的一个例子。
2.2 Transformer层
Transformer层通过MyModelLayers_类实现:
class MyModelLayers_(nn.Module):
def __init__(self, model, layer_idx):
"""
参数:
model: 模型实例
layer_idx: 当前层的索引编号
"""
super().__init__()
# ...
def forward(self, hidden_states, **kwargs):
# ...
pass
该类主要用于定义模型中的Transformer层，包括自注意力层、前馈神经网络层等。
对于fa后端，需要根据代码中实际的模型结构，决定是否添加残差和dropout。
2.3 归一化层
归一化层通过MyModelPreNorm_类实现:
class MyModelPreNorm_(nn.Module):
def __init__(self, model):
"""
参数:
model: 模型实例
"""
super().__init__()
# ...
def forward(self, hidden_states, **kwargs):
# ...
pass
该类主要用于定义模型中输出层前的归一化层。
2.4 输出层
输出层通过MyModelCls_类实现:
class MyModelCls_(nn.Module):
def __init__(self, model):
"""
参数:
model: 模型实例
"""
super().__init__()
# ...
def forward(self, hidden_states, **kwargs):
# ...
pass
该类主要用于定义模型的输出层。
为了增强代码的健壮性，该函数还需要支持一些额外的特性：Megatron序列并行、Ulysses序列并行（fa不支持）、并行求loss（fa不支持）,这方面的详细例子可以参考gpt_hf和gpt_fa。
注意：当使用hf后端时，获取logits_parallel需要直接引用原模型的.weight变量，这一点在FSDP中是不允许的，因此可以单独将获取logits_parallel的代码放在一个单独的函数中，用MyModelLoss_来表示，gpt_hf中展示了这样的一个例子。
在实现这些层时，需要特别注意，Transformer层中相同种类的层的forward函数输入张量（kwargs除外）和输出张量的格式和大小相同，这是为了方便更新模型信息，以保证流水线并行的正确性。例如在gpt_hf中，Transformer层的forward函数输入张量和输出张量的格式和大小相同，都是hidden_states。
2.5 构造序列化模型
构造序列化模型通过construct_sequential_model函数实现:
def construct_sequential_model(model, config):
"""
将模型转换为序列化版本
参数:
model: 原始模型实例
config: 模型配置对象，TransformerConfig
返回:
转换后的序列化模型
"""
model_ = PipeSequential()
# ...
这个函数将模型转化为PipeSequential 格式，它是一个特殊的序列容器，专门用于流水线并行。开发者只需要把模型按照顺序顺次通过add_module方法添加到PipeSequential中即可。
注意：如果使用了MyModelLoss_，还需要给其增加reset_parameters方法，以保证模型可以正确初始化。
2.6 模型信息
模型信息通过MyModelModelInfo类实现:
class MyModelModelInfo(ModelInfo):
def __init__(self, config, args):
super(MyModelModelInfo, self).__init__()
# ...
self.set_layernums(layernum_list)
self.set_shapes(layer_shapes_list)
self.set_dtypes(layer_dtypes_list)
self.set_module_types(module_types)
在该类中，需要赋值四个变量：layernums、shapes、dtypes、module_types，分别表示每种不同类型的Transformer层数，每种类型层的输入输出张量形状、每种类型层输入输出张量的数据类型、模型每一层的模型名称。
对于layernums，需要赋值一个列表，列表中的每个元素表示每种类型Transformer层的数量，例如对于GPT，列表的长度为1，因为GPT只有一种Decoder层，但对于T5，列表的长度为2，因为T5同时包含Encoder和Decoder层，这两种层的结构是不同的。
对于shapes，需要赋值一个列表，列表中的每个元素表示每种类型Transformer层的输入输出张量形状，通常是一个大小为[x,y]的列表，x表示Transformer层的种类，y表示每层输入输出张量的数量，列表中的每个值存储的是输入输出张量的形状。
对于dtypes，需要赋值一个列表，列表中的每个元素表示每种类型Transformer层的输入输出张量的数据类型，通常是一个大小为[x,y]的列表，x表示Transformer层的种类，y表示每层输入输出张量的数量，列表中的每个值存储的是输入输出张量的数据类型。
对于module_types，需要赋值一个列表，列表中的每个元素顺次表示模型中每一层的名称。
3 混合并行实现
混合并行实现通过MyModelModel_hybrid_parallel.py文件实现，该文件是连接模型与Galvatron并行系统的桥梁，主要负责构建支持混合并行的模型实例。
该文件主要实现了四个函数：get_hybrid_parallel_configs，construct_hybrid_parallel_model，get_mymodel_config，mymodel_model_hp。
3.1 获取混合并行配置
get_hybrid_parallel_configs函数用于获取混合并行策略，其实现格式如下：
def get_hybrid_parallel_configs(model_config, training_args):
hybrid_parallel_configs = get_hybrid_parallel_configs_api(model_config, training_args, MyModelModelInfo)
return hybrid_parallel_configs
该函数不需要任何改动，通过调用Galvatron的get_hybrid_parallel_configs_api函数获取混合并行策略，并返回一个字典，字典中包含混合并行策略信息。
3.2 构建混合并行模型
construct_hybrid_parallel_model函数用于构建混合并行模型，其实现格式如下：
def construct_hybrid_parallel_model(model, model_config, training_args, hybrid_parallel_configs):
# ...
hp_model = construct_hybrid_parallel_model_api(...)
return hp_model
该函数通过调用Galvatron的construct_hybrid_parallel_model_api函数构建混合并行模型，并返回一个支持混合并行的模型实例。具体来说，该API函数具体需要的参数和格式如下：
def construct_hybrid_parallel_model_api(
model, # 原始模型实例
model_config, # 模型配置对象
training_args, # 训练参数
hybrid_parallel_configs, # 混合并行配置
model_info, # 模型信息类
construct_sequential_model, # 构建序列化模型的函数
construct_tensor_parallel_model, # 构建张量并行模型的函数
wrap_block_name=None, # 需要包装FSDP的模块名称列��
wrap_checkpoint_block_name=None, # 需要添加检查点的模块名称列表
wrap_other_block_name=None, # 需要包装FSDP的其他模块名称列表
tied_wte_attr_names=None, # 权重绑定的属性名称列表
layernorm_name = [], # 层归一化的名称列表
all_block_name = None, # 所有模块的名称列表
load_module_func = None, # 加载模块的函数
):
# ...
pass
在此，我们额外对一些可能感到疑惑的可选参数进行解释：
wrap_block_name：需要包装FSDP的Transfomer层模块类列表。wrap_checkpoint_block_name：需要添加检查点的模块名称列表，通常是Transformer层。wrap_other_block_name：需要包装FSDP的其他模块名称列表，通常是Transformer层以外的其它层，注意这里如果定义了多个Embedding类，需要将所有细粒度Embedding类都添加到列表中。tied_wte_attr_names：权重绑定的属性名称列表，部分模型Vocab Embedding层和输出层的参数是相同的，对于需要这种需求的模型，开发者需要将模型第一层和最后一层中如何访问Vocab Embedding层的方式告诉Galvatron，例如对于gpt_hf，GPTVocabEmbedding_类在Embedding层通过self.wte访问，而输出层在Cls层直接通过self访问即可，因此tied_wte_attr_names为['wte'，'']。layernorm_name：用于标识Galvatron在不同的层该如何访问Layernorm的名称列表（不需要完整名称，只需要知道后缀名词即可），例如对于gpt_hf，Layernorm在GPTAttention_tp和GPTMLP_tp类中通过self.LayerNorm访问，在GPTPreNorm_中通过self.ln访问，因此layernorm_name为['LayerNorm', 'ln']。all_block_name：所有模块的名称列表，通常是wrap_block_name和wrap_other_block_name的并集。load_module_func：加载模块的函数，通常是定义在MyModelModel_checkpoint.py文件中的load_MyModel_module函数。
注意：虽然wrap_block_name、wrap_checkpoint_block_name、wrap_other_block_name、all_block_name这些参数在construct_hybrid_parallel_model_api中是可选参数，但为了保证模型可以正确初始化，这些参数必须传入。
3.3 获取模型配置
get_mymodel_config函数用于获取模型配置，其实现格式如下：
def get_mymodel_config(args, overwrite_args=True):
config = config_from_meta(args.model_size)
config = set_model_config(config, args, overwrite_args)
if hasattr(args, 'local_rank') and args.local_rank == 0:
print(config)
return config
3.4 构建混合并行模型
mymodel_model_hp函数用于构建混合并行模型，其实现格式如下：
def mymodel_model_hp(config, args):
hybrid_parallel_configs = get_hybrid_parallel_configs(model_config=config, training_args=args)
if args.local_rank == 0:
print("Creating Model...")
mymodel_model = MyModelModel_huggingface(config)
model = construct_hybrid_parallel_model(
model=mymodel_model,
model_config=config,
training_args=args,
hybrid_parallel_configs=hybrid_parallel_configs
)
return model
注意这里MyModelModel_huggingface是直接通过transformers获取到的模型，而不是Galvatron的模型。在huggingface中选择模型时，需要选择包含输出层的模型。
4 模型检查点保存加载实现（Experimental, 支持hf）
模型检查点保存加载实现通过MyModelModel_checkpoint.py文件实现，该文件定义了模型的检查点保存和加载实现，包括检查点的保存和加载函数。
该文件需要实现save_MyModel_module和load_MyModel_module函数。用于实现模型检查点的保存和加载。
Galvatron是按层存储和加载模型检查点的，因此在实现时需要注意按层进行加载和存储。
llama_hf中展示了如何实现模型检查点的保存和加载。
辅助文件说明
1 模型配置文件
模型配置文件定义了模型的配置，包括模型的结构、参数量等。
1.1 模型配置存储文件
meta_configs/MyModel-{MODEL_SIZE}b.json：模型配置文件，用于存储模型配置信息。
1.2 模型配置处理文件
meta_configs/config_utils.py：该文件主要负责处理模型配置相关的功能，其主要包括三部分：
获取模型配置信息：通过调用
config_from_meta函数获取模型配置信息，并写入到TransformerConfig中。修改模型配置信息：通过调用
set_model_config函数，根据传入的arguments修改模型配置信息，并通过overwrite_megatron_args和overwrite_model_args函数修改arguments中的模型配置信息。获取模型相关信息：通过
model_name函数获取模型名称，通过model_layer_configs函数获取模型每一层的配置信息。
2 训练文件
训练文件主要定义了训练相关的功能，包括数据加载、模型训练等。
2.1 训练主文件
train_dist.py：该文件主要负责分布式训练相关的功能。
一个完整的示例如下：
def train(args):
# 初始化分布式训练环境
local_rank = args.local_rank
rank = torch.distributed.get_rank()
torch.cuda.set_device(local_rank)
device = torch.device("cuda", local_rank)
world_size = torch.distributed.get_world_size()
config = get_mymodel_config(args)
model = mymodel_model_hp(config, args)
# 创建数据集
if local_rank == 0:
print("Creating Dataset...")
# 设置数据集相关参数
set_megatron_args_for_dataset(args, model,
model.sp_groups_whole[0] if args.vocab_sp else model.tp_groups_whole[0],
model.dp_groups_whole[0])
if local_rank == 0:
_print_args("arguments", args)
# 获取数据迭代器
train_data_iterator, valid_data_iterator, test_data_iterator = get_train_valid_test_data_iterators()
# 创建优化器和学习率调度器
optimizer, opt_param_scheduler = get_optimizer_and_param_scheduler(model, args)
# 设置性能分析器
path = os.path.dirname(os.path.abspath(__file__))
profiler = GalvatronProfiler(args)
profiler.set_profiler_dist(path, model_layer_configs(config), model_name(config), start_iter=0)
# 记录模型创建后的内存使用情况
profiler.profile_memory(0, "After creating model")
if local_rank == 0:
print("Start training...")
# 训练循环
for iter in range(args.iteration, args.train_iters):
# 获取一个批次的数据
tokens, kwargs, loss_func = get_batch(train_data_iterator)
# 记录开始时间和内存使用
profiler.profile_time_start(iter)
profiler.profile_memory(iter, "Before Forward")
# 准备输入数据
input_ids = tokens
batch = [input_ids]
# 前向传播和反向传播
loss = model.forward_backward(batch, iter, profiler,
loss_func=loss_func,
**kwargs)
# 记录反向传播后的内存使用
profiler.profile_memory(iter, "After Backward")
# 梯度裁剪
total_norm = clip_grad_norm(model, args.clip_grad)
# 优化器步骤
optimizer.step()
# 学习率调度器步骤
opt_param_scheduler.step(increment=args.global_batch_size)
# 记录优化器步骤后的内存使用
profiler.profile_memory(iter, "After optimizer_step")
# 清零梯度
optimizer.zero_grad()
# 更新性能统计信息
profiler.post_profile_memory(iter)
# 获取当前学习率
for param_group in optimizer.param_groups:
learning_rate = param_group['lr']
# 记录本次迭代的性能指标
profiler.profile_time_end(iter, loss, learning_rate, total_norm)
# 同步所有进程
torch.distributed.barrier()
# 定期保存模型检查点
if args.save != None and (iter + 1) % args.save_interval == 0:
save_llama_module(args.save, model, optimizer, opt_param_scheduler, iter + 1, args)
if __name__ == '__main__':
# 初始化Galvatron训练环境
args = initialize_galvatron(model_args, mode='train_dist')
# 设置随机种子以确保可重复性
set_seed()
# 开始训练
train(args)
train_dist_random.py：该文件主要负责分布式训练相关的功能，与
train_dist.py类似，但使用随机数据进行训练。
2.2 数据加载文件
dataloader.py：该文件主要负责数据加载相关的功能，其主要包括两部分：
随机数据加载：创建生成随机token的dataset，并创建collate_fn函数，将随机token转换为模型输入。 如下是一个随机数据加载的示例：
def random_get_ltor_masks_and_position_ids(data): """Build masks and position id for left to right model.""" micro_batch_size, seq_length = data.size() att_mask_batch = 1 attention_mask = torch.tril(torch.ones( (att_mask_batch, seq_length, seq_length), device=data.device)).view( att_mask_batch, 1, seq_length, seq_length) attention_mask = (attention_mask < 0.5) return attention_mask def random_collate_fn(batch): # 将batch中的数据堆叠，并返回对应格式的数据 tokens_ = torch.stack(batch, dim=0) labels = tokens_[:, 1:].contiguous() tokens = tokens_[:, :-1].contiguous() args = get_args() if not args.use_flash_attn: attention_mask = random_get_ltor_masks_and_position_ids(tokens) else: attention_mask = None return tokens, {"attention_mask":attention_mask, "labels" : labels}, None class DataLoaderForMyModel(Dataset): def __init__(self, args, device, dataset_size = 2560 * 16): self.vocab_size = args.vocab_size self.sentence_length = args.seq_length self.dataset_size = dataset_size # 随机生成每个样本的实际长度（1到最大长度之间） self.data_length = np.random.randint(1,self.sentence_length+1,(self.dataset_size,)) self.device = device # 生成随机输入数据 self.input_ids = [] for i in range(self.dataset_size): sentence = np.random.randint(0,self.vocab_size,(self.sentence_length,)) sentence[self.data_length[i]:] = 0 mask = np.ones((self.sentence_length,)) mask[self.data_length[i]:] = 0 padding_sentence = np.zeros(self.sentence_length + 1, dtype=sentence.dtype) padding_sentence[:self.sentence_length] = sentence self.input_ids.append(padding_sentence) self.input_ids = np.array(self.input_ids) def __len__(self): return self.dataset_size def __getitem__(self, idx): if idx >= self.dataset_size: raise IndexError input_ids = torch.LongTensor(self.input_ids[idx]).to(self.device) return input_ids
具体的trainloader由以下代码创建：
trainloader = distributed_dataloader( dataset=DataLoaderForGPT(args, device), global_bsz=args.global_train_batch_size, shuffle=True, args=args, group = model.dp_groups_whole[0].group, collate_fn = random_collate_fn )
其中
distributed_dataloader函数是Galvatron提供的分布式数据加载器，用于创建分布式数据加载器。真实数据加载：创建真实数据加载器，并设计loss计算函数。
真实数据加载的实现基于Megatron dataset，主要包含
train_valid_test_datasets_provider、get_train_valid_test_data_iterators、get_batch、loss_func等函数。一个具体实现的例子可以参考gpt_hf。主要注意的是，
get_batch函数返回一个tuple，tuple中包含三个元素，分别是：输入数据：通常是一个token序列，torch.Tensor类型。
其他输入数据：通常是字典类型，包含position_ids、attention_mask、labels等。
loss计算函数：通过调用
loss_func(output_tensor)函数可以直接计算出loss。
注意：这里的输入数据要和
MyModelModel_sequential.py文件中Embedding层的输入数据格式保持一致。而其他数据则作为**kwargs在模型层之间传递。
2.3 性能分析文件
profiler.py：该文件主要负责性能分析相关的功能，其内容如下：
if __name__ == '__main__':
# 初始化Galvatron性能分析环境
args = initialize_galvatron(model_args, mode='profile')
# 加载模型配置
config = get_mymodel_config(args, overwrite_args=False)
# 创建性能分析器实例
profiler = GalvatronProfiler(args)
# 获取当前文件的目录路径
path = os.path.dirname(os.path.abspath(__file__))
# 设置性能分析器启动器
profiler.set_profiler_launcher(path, layernum_arg_names(), model_name(config))
# 启动性能分析脚本
profiler.launch_profiling_scripts()
# 处理收集到的性能数据
profiler.process_profiled_data()
2.4 策略搜索文件
search_dist.py：该文件主要负责策略搜索相关的功能，其内容如下：
if __name__ == '__main__':
args = initialize_galvatron(model_args, mode='search')
config = get_mymodel_config(args, overwrite_args=True)
path = os.path.dirname(os.path.abspath(__file__))
print(args)
print(config)
# 创建策略搜索引擎实例
search_engine = GalvatronSearchEngine(args)
# 设置搜索引擎的基本信息
search_engine.set_search_engine_info(path, model_layer_configs(config), model_name(config))
# 初始化搜索引擎
search_engine.initialize_search_engine()
# 进行策略搜索
search_engine.parallelism_optimization()
3 脚本文件
scripst文件夹中主要包含一些脚本文件，用于实现模型训练、性能分析、策略搜索等功能。
主要包含五种不同的脚本：
profile_computation.sh：用于性能分析，计算模型在不同配置下的计算性能。
profile_memory.sh：用于性能分析，计算模型在不同配置下的内存使用情况。
search_dist.sh：用于策略搜索，搜索模型在不同配置下的最优策略。
train_dist.sh：用于模型训练，训练模型。
train_dist_random.sh：用于模型训练，使用随机数据训练模型。