# 全局注册表
MODEL_ARCH_CONFIGS = {}          # 模型架构配置
MODEL_ARCH_TO_FAMILY = {}        # 架构到模型族映射
MODEL_FAMILY_TO_ARCHS = {}       # 模型族到架构映射
MODEL_FAMILY_TO_PROVIDER = {}    # 模型族到提供者映射

# 注册模型配置
@register_model_config(model_family="llava-ov-1.5", model_arch="llava-ov-1.5-4b")
def llava_one_vision_1_5_4b():
    return LlavaOnevision1_5Config(
        num_layers=36,
        hidden_size=2560,
        ffn_hidden_size=9728,
        num_attention_heads=32,
        group_query_attention=True,
        num_query_groups=8,
        vocab_size_in_config_file=151936,
        # ... 更多配置
    )

# 注册模型提供者
@register_model_provider(model_family=[VisionLanguageModelFamilies.LLAVA_OV_1_5])
def rice_vl_model_provider(pre_process=True, post_process=True):
    # 模型构建逻辑
    pass

@register_model_trainer(
    model_family=[constants.VisionLanguageModelFamilies.LLAVA_OV_1_5],
    training_phase=constants.TrainingPhase.PRETRAIN
)
def default_pretrain_trainer(train_args):
    trainer = MegatronTrainer(
        train_args=train_args,
        train_valid_test_dataset_provider=train_valid_test_dataset_provider,
        model_provider=model_provider,
        model_type=model_type,
        forward_step_func=forward_step,
    )
    return trainer

class LlavaOnevision1_5(MegatronModule):
    def __init__(self, ...):
        # 1. 视觉编码器（Vision Encoder）
        if self.add_encoder:
            self.vision_model = RiceViTModel(
                vision_config,
                vision_layer_spec,
            )
            
        # 2. 适配器（Adapter）- 连接视觉和语言模态
            self.adapter = Adapter(
                adapter_config,
                adapter_layer_spec,
                input_size=vision_config.hidden_size,
                output_size=language_config.hidden_size,
            )
        
        # 3. 语言模型（Language Model）
        if self.add_decoder:
            self.language_model = QwenModel(
                config=language_config,
                transformer_layer_spec=language_layer_spec,
                vocab_size=language_vocab_size,
                max_sequence_length=language_max_sequence_length,
                # ... 更多参数
            )

def forward(
    self,
    images: torch.Tensor,              # 图像输入
    image_grid_thw: torch.Tensor,      # 图像网格维度（时间、高、宽）
    input_ids: torch.Tensor,           # 文本token ID
    position_ids: torch.Tensor,        # 位置ID
    attention_mask: torch.Tensor,      # 注意力掩码
    labels: torch.Tensor = None,       # 标签（训练时使用）
    packed_seq_params: PackedSeqParams = None,  # 序列打包参数
    inference_params: InferenceParams = None,    # 推理参数
):
    # 步骤1：视觉编码
    if self.add_encoder and images is not None:
        # 通过视觉模型提取图像特征
        image_embeddings, window_index = self.vision_model(
            images, 
            grid_thw=image_grid_thw
        )
        # 通过适配器映射到语言模型维度
        image_embeddings = self.adapter(image_embeddings, window_index)
    
    # 步骤2：文本编码
    if self.pre_process:
        language_embeddings = self.language_model.embedding(
            input_ids=input_ids, 
            position_ids=None
        )
        
        # 步骤3：融合视觉和语言嵌入
        if images is not None:
            # 找到图像token位置并替换为视觉嵌入
            images_mask = (input_ids == self.config.image_token_id)
            combined_embeddings = language_embeddings.masked_scatter(
                images_mask, 
                image_embeddings
            )
    
    # 步骤4：通过语言模型生成输出
    output = self.language_model(
        input_ids=None,
        decoder_input=combined_embeddings,
        attention_mask=attention_mask,
        labels=labels,
        # ... 其他参数
    )
    
    return output

class RiceViTModel(VisionModel):
    def __init__(self, config, transformer_layer_spec, spatial_merge_size=2):
        # 1. Patch嵌入层
        self.patch_embed = PatchEmbed(
            patch_size=config.patch_size,  # 14
            in_channels=config.in_channels,  # 3
            embed_dim=config.hidden_size,    # 1024
        )
        
        # 2. 旋转位置编码
        self.rotary_pos_emb = VisionRotaryEmbedding(
            dim=config.kv_channels // 2
        )
        
        # 3. Transformer块
        self.decoder = TransformerBlock(
            config=config,
            spec=transformer_layer_spec,
        )
        
        # 4. 类别token（用于全局特征）
        self.class_embedding = torch.randn(config.hidden_size)
        self.class_pos_emb = torch.randn(1, config.kv_channels // 2)
        
        # 5. 预层归一化
        self.pre_layernorm = torch.nn.LayerNorm(
            config.hidden_size, 
            eps=1e-4
        )

def rot_pos_emb(self, grid_thw):
    """
    为每个图像patch生成2D旋转位置编码
    grid_thw: [num_images, 3] - (时间, 高度, 宽度)
    """
    pos_ids = []
    for t, h, w in grid_thw:
        # 生成高度位置ID
        hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w)
        hpos_ids = hpos_ids.reshape(
            h // self.spatial_merge_size,
            self.spatial_merge_size,
            w // self.spatial_merge_size,
            self.spatial_merge_size,
        )
        hpos_ids = hpos_ids.permute(0, 2, 1, 3).flatten()
        
        # 生成宽度位置ID
        wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1)
        # ... 类似处理 ...
        
        pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1))
    
    # 应用旋转位置编码
    rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size)
    rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1)
    
    return rotary_pos_emb

class Adapter(MegatronModule):
    """
    2层MLP适配器
    Input: [seq_len, hidden_size_vision]
    Output: [seq_len, hidden_size_language]
    """
    def __init__(self, config, layer_spec, input_size, output_size):
        # 第一层：投影到中间维度
        self.linear_fc1 = build_module(
            layer_spec.linear_fc1,
            input_size,
            config.ffn_hidden_size,
        )
        
        # 激活函数（通常是GELU）
        self.activation_func = config.activation_func
        
        # 第二层：投影到输出维度
        self.linear_fc2 = build_module(
            layer_spec.linear_fc2,
            config.ffn_hidden_size,
            output_size,
        )

VisionConfig(
    num_layers=24,                    # Transformer层数
    hidden_size=1024,                 # 隐藏层维度
    ffn_hidden_size=4096,            # FFN维度
    num_attention_heads=16,          # 注意力头数
    patch_size=14,                   # Patch大小
    image_size=(1344, 1344),         # 最大图像尺寸
    kv_channels=64,                  # KV通道数
    normalization="LayerNorm",       # 归一化类型
    position_embedding_type="rope"   # 位置编码类型
)

class MegatronTrainer(object):
    """Megatron训练器基类"""
    def __init__(
        self,
        train_args,                              # 训练参数
        train_valid_test_dataset_provider,       # 数据集提供者
        model_provider,                          # 模型提供者
        model_type,                              # 模型类型
        forward_step_func,                       # 前向传播函数
        process_non_loss_data_func=None,         # 非损失数据处理函数
        get_embedding_ranks=None,                # 获取嵌入层所在rank
        get_position_embedding_ranks=None,       # 获取位置嵌入所在rank
    ):
        self.train_args = train_args
        self.train_valid_test_datasets_provider = train_valid_test_dataset_provider
        self.model_provider = model_provider
        self.model_type = model_type
        self.forward_step_func = forward_step_func
        # ...
    
    def train(self):
        """启动训练"""
        pretrain(
            train_args=self.train_args,
            train_valid_test_dataset_provider=self.train_valid_test_datasets_provider,
            model_provider=self.model_provider,
            model_type=self.model_type,
            forward_step_func=self.forward_step_func,
            # ...
        )

def pretrain(
    train_args,
    train_valid_test_dataset_provider,
    model_provider,
    model_type,
    forward_step_func,
    process_non_loss_data_func=None,
):
    """
    主训练流程：
    1. 初始化Megatron
    2. 构建模型、优化器和学习率调度器
    3. 获取数据迭代器
    4. 执行训练循环
    """
    
    # 步骤1：初始化
    initialize_aiak_megatron(
        args=train_args,
        get_embedding_ranks=get_embedding_ranks,
        get_position_embedding_ranks=get_position_embedding_ranks
    )
    
    # 步骤2：构建模型和优化器
    model, ema, optimizer, opt_param_scheduler = setup_model_and_optimizer(
        model_provider, 
        model_type, 
        checkpointing_context=checkpointing_context
    )
    
    # 步骤3：构建数据迭代器
    train_data_iterator, valid_data_iterator, test_data_iterator = \
        build_train_valid_test_data_iterators(
            train_valid_test_dataset_provider
        )
    
    # 步骤4：训练循环
    if args.do_train and args.train_iters > 0:
        iteration, num_floating_point_operations_so_far = train(
            forward_step_func=forward_step_func,
            model=model,
            ema=ema,
            optimizer=optimizer,
            opt_param_scheduler=opt_param_scheduler,
            train_data_iterator=train_data_iterator,
            valid_data_iterator=valid_data_iterator,
            # ...
        )

def train(forward_step_func, model, ema, optimizer, ...):
    """主训练循环"""
    
    while iteration < args.train_iters:
        # 1. 执行训练步骤
        loss_dict, skipped_iter, should_checkpoint, should_exit, exit_code, \
        grad_norm, num_zeros_in_grad = train_step(
            forward_step_func,
            train_data_iterator,
            model,
            optimizer,
            opt_param_scheduler,
            config
        )
        
        # 2. 更新EMA模型（如果启用）
        if args.enable_ema and ema is not None:
            for e, m in zip(ema, model):
                update_ema(e, m, rate=args.ema_decay)
        
        # 3. 日志记录
        report_memory_flag = training_log(
            loss_dict, 
            total_loss_dict,
            learning_rate,
            iteration, 
            loss_scale,
            # ...
        )
        
        # 4. 验证评估（按间隔）
        if args.eval_interval and iteration % args.eval_interval == 0:
            evaluate_and_print_results(
                prefix, 
                forward_step_func,
                valid_data_iterator, 
                model,
                iteration, 
                # ...
            )
        
        # 5. 保存检查点
        should_exit = checkpoint_and_decide_exit(
            model, 
            optimizer, 
            opt_param_scheduler, 
            iteration,
            num_floating_point_operations_so_far,
            # ...
        )
        
        iteration += 1

def train_step(forward_step_func, data_iterator, model, optimizer, ...):
    """单个训练步骤"""
    
    # 1. 清空梯度
    for model_chunk in model:
        model_chunk.zero_grad_buffer()
    optimizer.zero_grad()
    
    # 2. 前向和反向传播
    forward_backward_func = get_forward_backward_func()
    losses_reduced = forward_backward_func(
        forward_step_func=forward_step_func,
        data_iterator=data_iterator,
        model=model,
        num_microbatches=get_num_microbatches(),
        seq_length=args.seq_length,
        micro_batch_size=args.micro_batch_size,
        forward_only=False
    )
    
    # 3. 更新参数
    update_successful, grad_norm, num_zeros_in_grad = optimizer.step()
    
    # 4. 更新学习率
    if update_successful:
        increment = get_num_microbatches() * args.micro_batch_size * args.data_parallel_size
        opt_param_scheduler.step(increment=increment)
        skipped_iter = 0
    else:
        skipped_iter = 1
    
    # 5. 计算平均损失
    loss_reduced = {}
    for key in losses_reduced[0].keys():
        numerator = sum(x[key][0] for x in losses_reduced)
        denominator = sum(x[key][1] for x in losses_reduced)
        loss_reduced[key] = numerator / denominator
    
    return loss_reduced, skipped_iter, ...

def forward_step(data_iterator, model):
    """预训练前向传播步骤"""
    
    # 1. 获取批次数据
    images, image_grid_thw, pixel_values_videos, video_grid_thw, \
    input_ids, position_ids, attention_mask, \
    labels, loss_mask, attn_mask_type, packed_seq_params = get_batch(data_iterator)
    
    # 2. 模型前向传播
    output_tensor = model(
        images,                    # 图像输入
        image_grid_thw,           # 图像网格
        input_ids,                # 文本token
        position_ids,             # 位置ID
        attention_mask,           # 注意力掩码
        attn_mask_type,           # 掩码类型
        labels,                   # 标签
        packed_seq_params,        # 打包序列参数
        pixel_values_videos=pixel_values_videos,  # 视频输入
        video_grid_thw=video_grid_thw             # 视频网格
    )
    
    # 3. 返回输出和损失函数
    return output_tensor, partial(loss_func, loss_mask)

def loss_func(loss_mask: torch.Tensor, output_tensor: torch.Tensor):
    """
    损失函数计算
    
    Args:
        loss_mask: 损失掩码，标记哪些位置需要计算损失
        output_tensor: 模型输出的损失张量
    
    Returns:
        loss: 标量损失值
        local_num_tokens: 本地token数量
        loss_reduced_dict: 包含各种统计信息的字典
    """
    
    losses = output_tensor.float()
    loss_mask = loss_mask.view(-1).float()
    
    # 计算加权损失
    total_tokens = loss_mask.sum()
    loss = torch.cat([
        torch.sum(losses.view(-1) * loss_mask).view(1), 
        total_tokens.view(1)
    ])
    
    # Context Parallel归约
    if args.context_parallel_size > 1:
        torch.distributed.all_reduce(loss, group=mpu.get_context_parallel_group())
    
    # NaN检查
    if args.check_for_nan_in_loss_and_grad:
        assert not loss[0].isnan(), f'Found NaN in loss on rank {rank}'
    
    # 跨数据并行组归约
    reporting_loss = loss.clone().detach()
    torch.distributed.all_reduce(reporting_loss, group=mpu.get_data_parallel_group())
    
    loss_reduced_dict = {
        'lm loss': (reporting_loss[0], reporting_loss[1])
    }
    
    # 可变序列长度处理
    if args.variable_seq_lengths:
        num_input_tokens = output_tensor.shape[0] * output_tensor.shape[1]
        input_tokens = torch.tensor(num_input_tokens, dtype=torch.int)
        torch.distributed.all_reduce(input_tokens, group=mpu.get_data_parallel_group())
        loss_reduced_dict["total_inputs"] = input_tokens.item()
    
    return (
        loss[0] * args.context_parallel_size,
        loss[1].to(torch.int),
        loss_reduced_dict
    )

def setup_model_and_optimizer(model_provider_func, model_type):
    """设置模型和优化器"""
    
    # 1. 构建模型
    model = get_model(model_provider_func, model_type)
    
    # 2. 构建优化器配置
    config = OptimizerConfig(
        # 从args中提取配置
        lr=args.lr,
        weight_decay=args.weight_decay,
        adam_beta1=args.adam_beta1,
        adam_beta2=args.adam_beta2,
        # ...
    )
    
    # 3. 创建优化器
    optimizer = get_megatron_optimizer(
        config, 
        model,
        use_gloo_process_groups=args.enable_gloo_process_groups
    )
    
    # 4. 创建学习率调度器
    opt_param_scheduler = get_optimizer_param_scheduler(optimizer)
    
    # 5. 加载检查点（如果存在）
    if args.load is not None:
        args.iteration, args.num_floating_point_operations_so_far = \
            load_checkpoint(model, optimizer, opt_param_scheduler)
    else:
        args.iteration = 0
    
    # 6. 创建EMA模型（如果启用）
    if args.enable_ema:
        ema = get_model(model_provider_func, model_type)
        if args.iteration == 0:
            # 初始化EMA为模型参数
            for e, m in zip(ema, model):
                update_ema(e, m, rate=0)
        else:
            load_checkpoint(ema, None, None, load_arg="load_ema")
    
    return model, ema, optimizer, opt_param_scheduler

@dataclass
class SFTDatasetConfig(BlendedHuggingFaceDatasetConfig):
    """SFT数据集配置"""
    
    # 数据路径和混合
    dataset_config_file: str          # 数据集配置文件路径
    blend: Optional[List]              # 数据混合比例
    blend_per_split: Optional[List]    # 各分割的混合比例
    
    # 模板和处理器
    chat_template: Optional[ChatTemplate]    # 对话模板
    processor: Optional[ProcessorMixin]      # 数据处理器
    
    # 训练配置
    sequence_length: int                     # 最大序列长度
    ignore_index: int = -100                 # 忽略索引
    train_on_prompt: bool = False           # 是否在提示上训练
    
    # 数据增强
    packing: bool = False                   # 是否打包数据
    sort_batch: bool = False                # 是否排序批次
    packing_batch_size: int = 10000        # 打包批次大小
    
    # 流式加载
    streaming: bool = False                 # 是否流式加载
    streaming_buffer_size: int = 16384     # 流式缓冲区大小
    
    # 并行配置
    context_parallel_size: Optional[int] = None  # 上下文并行大小

class SFTDataFormats:
    """支持的SFT数据格式"""
    ALPACA = "alpaca"           # Alpaca格式
    SHAREGPT = "sharegpt"       # ShareGPT格式
    OPENAI = "openai"           # OpenAI格式

class SFTDataset(HuggingFaceDataset):
    """SFT数据集类"""
    
    def __init__(self, config: SFTDatasetConfig):
        super().__init__(config)
        
        # 1. 加载原始数据集
        raw_datasets = self._load_datasets()
        
        # 2. 格式转换
        processed_datasets = self._convert_format(raw_datasets)
        
        # 3. 分词处理
        tokenized_datasets = self._tokenize(processed_datasets)
        
        # 4. 数据打包（可选）
        if config.packing:
            tokenized_datasets = self._pack_samples(tokenized_datasets)
        
        # 5. 排序（可选）
        if config.sort_batch:
            tokenized_datasets = self._sort_by_length(tokenized_datasets)
        
        self.dataset = tokenized_datasets

class Qwen2VLTaskEncoder:
    """Qwen2-VL任务编码器"""
    
    def __init__(self, args):
        self.args = args
        self.tokenizer = get_tokenizer()
        self.processor = AutoProcessor.from_pretrained(
            args.hf_tokenizer_path,
            trust_remote_code=True
        )
    
    def encode_sample(self, sample):
        """
        编码单个样本
        
        Args:
            sample: 包含文本和图像的字典
                {
                    "messages": [...],        # 对话消息
                    "images": [...],          # 图像路径列表
                    "videos": [...],          # 视频路径列表（可选）
                }
        
        Returns:
            编码后的样本，包含：
                - tokens: token ID
                - labels: 标签
                - imgs: 图像张量
                - image_grid_thw: 图像网格维度
                - attn_mask: 注意力掩码
                - cu_lengths: 累积长度（用于打包）
        """
        
        # 1. 应用对话模板
        text = self.processor.apply_chat_template(
            sample["messages"],
            tokenize=False,
            add_generation_prompt=False
        )
        
        # 2. 处理图像
        image_inputs = None
        if sample.get("images"):
            images = [Image.open(img_path) for img_path in sample["images"]]
            image_inputs = self.processor.image_processor(
                images,
                return_tensors="pt"
            )
        
        # 3. 分词
        text_inputs = self.tokenizer(
            text,
            return_tensors="pt",
            padding=False,
            truncation=True,
            max_length=self.args.seq_length
        )
        
        # 4. 合并输入
        encoded = {
            "tokens": text_inputs["input_ids"],
            "labels": self._build_labels(text_inputs, sample),
            "attn_mask": text_inputs["attention_mask"],
        }
        
        if image_inputs:
            encoded.update({
                "imgs": image_inputs["pixel_values"],
                "image_grid_thw": image_inputs["image_grid_thw"],
            })
        
        return encoded

def pack_samples(samples, max_length, context_parallel_size=1):
    """
    将多个样本打包到固定长度
    
    Args:
        samples: 样本列表
        max_length: 最大序列长度
        context_parallel_size: 上下文并行大小
    
    Returns:
        打包后的样本
    """
    
    packed_samples = []
    current_pack = {
        "input_ids": [],
        "labels": [],
        "attention_mask": [],
        "cu_seqlens": [0],  # 累积序列长度
        "images": [],
        "image_grid_thw": [],
    }
    current_length = 0
    
    for sample in samples:
        sample_length = len(sample["input_ids"])
        
        # 检查是否可以添加到当前pack
        if current_length + sample_length <= max_length:
            # 添加到当前pack
            current_pack["input_ids"].extend(sample["input_ids"])
            current_pack["labels"].extend(sample["labels"])
            current_pack["attention_mask"].extend(sample["attention_mask"])
            current_pack["cu_seqlens"].append(current_length + sample_length)
            
            if "images" in sample:
                current_pack["images"].append(sample["images"])
                current_pack["image_grid_thw"].append(sample["image_grid_thw"])
            
            current_length += sample_length
        else:
            # 当前pack已满，开始新pack
            if current_length > 0:
                # 填充到max_length（考虑CP）
                pad_length = max_length - current_length
                if context_parallel_size > 1:
                    # 确保长度可被CP整除
                    pad_length = (max_length // context_parallel_size) * context_parallel_size - current_length
                
                current_pack["input_ids"].extend([0] * pad_length)
                current_pack["labels"].extend([-100] * pad_length)
                current_pack["attention_mask"].extend([0] * pad_length)
                
                packed_samples.append(current_pack)
            
            # 开始新pack
            current_pack = initialize_new_pack()
            current_length = 0
    
    return packed_samples

def get_train_loader(dataset, collator):
    """
    创建训练数据加载器
    
    Args:
        dataset: 数据集对象
        collator: 数据整理函数
    
    Returns:
        数据加载器迭代器
    """
    args = get_args()
    
    # 1. 创建数据采样器
    if args.dataloader_type == 'cyclic':
        # 循环数据加载器（用于WebDataset）
        dataloader = build_pretraining_data_loader(dataset)
    else:
        # 标准数据加载器
        from torch.utils.data import DataLoader
        
        dataloader = DataLoader(
            dataset,
            batch_size=args.micro_batch_size,
            shuffle=True,
            num_workers=args.num_workers,
            collate_fn=collator,
            pin_memory=True,
            drop_last=True,
        )
    
    # 2. 创建循环迭代器
    def cyclic_iter(dataloader):
        while True:
            for batch in dataloader:
                yield batch
    
    return cyclic_iter(dataloader)

class MultiModalDataCollatorForSupervisedDataset:
    """多模态监督数据集整理器"""
    
    def __init__(self, processor, plugin):
        self.processor = processor
        self.plugin = plugin
    
    def __call__(self, samples):
        """
        整理一个批次的样本
        
        Args:
            samples: 样本列表
        
        Returns:
            批次字典
        """
        
        # 1. 提取各个字段
        input_ids = [sample["input_ids"] for sample in samples]
        labels = [sample["labels"] for sample in samples]
        attention_masks = [sample["attention_mask"] for sample in samples]
        
        # 2. 填充到相同长度
        max_length = max(len(ids) for ids in input_ids)
        
        padded_input_ids = []
        padded_labels = []
        padded_attention_masks = []
        
        for ids, labs, mask in zip(input_ids, labels, attention_masks):
            pad_length = max_length - len(ids)
            
            padded_input_ids.append(ids + [0] * pad_length)
            padded_labels.append(labs + [-100] * pad_length)
            padded_attention_masks.append(mask + [0] * pad_length)
        
        # 3. 转换为张量
        batch = {
            "input_ids": torch.tensor(padded_input_ids, dtype=torch.long),
            "labels": torch.tensor(padded_labels, dtype=torch.long),
            "attention_mask": torch.tensor(padded_attention_masks, dtype=torch.long),
        }
        
        # 4. 处理图像（如果存在）
        if "images" in samples[0]:
            images = []
            image_grid_thw = []
            
            for sample in samples:
                if sample["images"] is not None:
                    images.append(sample["images"])
                    image_grid_thw.append(sample["image_grid_thw"])
            
            if images:
                batch["images"] = torch.cat(images, dim=0)
                batch["image_grid_thw"] = torch.tensor(image_grid_thw, dtype=torch.long)
        
        return batch

# examples/llava_ov_1_5/quick_start/stage_1_alignment_llava_ov_4b.sh

MODEL_ARGS=(
    --model-name llava-ov-1.5-4b
)

TRAINING_ARGS=(
    --training-phase sft
    --trainable-modules adapter        # 只训练adapter
    --seq-length 32768
    --micro-batch-size 1
    --global-batch-size 8
    --lr 1.0e-4
    --train-iters 2500
    --bf16
)

def main():
    """训练主入口"""
    
    # 1. 解析参数
    args = parse_train_args()
    
    # 2. 构建训练器
    trainer = build_model_trainer(args)
    
    # 3. 启动训练
    trainer.train()

def parse_train_args():
    """解析训练参数"""
    
    args = parse_arguments(
        extra_args_provider=aiak_extra_train_args_provider,
        validate_extra_args_provider=validate_aiak_extra_args,
    )
    
    return args

def aiak_extra_train_args_provider(parser):
    """添加AIAK自定义参数"""
    
    # 1. 模型参数
    parser = _add_extra_model_args(parser)
    
    # 2. 分词器参数
    parser = _add_extra_tokenizer_args(parser)
    
    # 3. SFT参数
    parser = _add_extra_sft_args(parser)
    
    # 4. 多模态参数
    parser = _add_extra_multimodal_args(parser)
    
    # 5. 并行参数
    parser = _add_extra_parallel_args(parser)
    
    return parser

--model-name llava-ov-1.5-4b          # 模型架构
--trainable-modules adapter           # 可训练模块（all/adapter/language_model/vision_model）

--data-path /path/to/dataset          # 数据路径
--dataloader-type external            # 数据加载器类型
--split 100,0,0                       # 训练/验证/测试分割
--chat-template qwen2-vl              # 对话模板
--num-workers 16                      # 数据加载工作线程

--training-phase sft                  # 训练阶段（pretrain/sft）
--seq-length 32768                    # 序列长度
--micro-batch-size 1                  # 微批次大小
--global-batch-size 8                 # 全局批次大小
--train-iters 2500                    # 训练迭代次数

--lr 1.0e-4                           # 学习率
--min-lr 1.0e-6                       # 最小学习率
--lr-decay-style cosine               # 学习率衰减策略
--lr-warmup-fraction 0.002            # 学习率预热比例
--clip-grad 1.0                       # 梯度裁剪

--optimizer adam                      # 优化器
--adam-beta1 0.9                      # Adam beta1
--adam-beta2 0.99                     # Adam beta2
--weight-decay 0                      # 权重衰减

--tensor-model-parallel-size 1        # 张量并行度
--pipeline-model-parallel-size 1      # 流水线并行度
--use-distributed-optimizer           # 使用分布式优化器
--attention-backend flash             # 注意力后端（flash/native）

--bf16                                # 使用BF16精度
--recompute-granularity full          # 重计算粒度
--recompute-method uniform            # 重计算方法
--recompute-num-layers 4              # 重计算层数

# 8GPU单节点，使用TP=2, PP=1
GPUS_PER_NODE=8
TP=2
PP=1

# 实际数据并行度 = 8 / (2 * 1) = 4
# 全局批次大小 = 微批次大小 * 数据并行度 * 梯度累积步数

TRAINING_ARGS=(
    --recompute-granularity full      # 完全重计算
    --recompute-method uniform        # 均匀选择层
    --recompute-num-layers 4          # 重计算4层
)

MODEL_PARALLEL_ARGS=(
    --use-distributed-optimizer       # 启用分布式优化器
)

MODEL_PARALLEL_ARGS=(
    --attention-backend flash         # 使用Flash Attention
)

TRAINING_ARGS=(
    --save /path/to/checkpoint        # 检查点保存路径
    --save-interval 2000              # 保存间隔
    --ckpt-format torch               # 检查点格式（torch/torch_dist/zarr）
)

TRAINING_ARGS=(
    --load /path/to/checkpoint        # 检查点加载路径
    --ckpt-fully-parallel-load        # 完全并行加载
)

TRAINING_ARGS=(
    --async-save                      # 启用异步保存
)

#!/bin/bash
LOAD=$1    # HuggingFace模型路径
SAVE=$2    # Megatron输出路径
TP=$3      # 张量并行度
PP=$4      # 流水线并行度

# 步骤1：转换语言模型
python $CONVERT_CHECKPOINT_PATH/model.py \
    --load_platform=huggingface \
    --save_platform=mcore \
    --common_config_path=$CONVERT_CHECKPOINT_PATH/config/llava-ov-1.5-4b/qwen3.json \
    --tensor_model_parallel_size=$TP \
    --pipeline_model_parallel_size=$PP \
    --load_ckpt_path=$LOAD \
    --save_ckpt_path=$SAVE_LANGUAGE_MODEL

# 步骤2：转换视觉模型
python $CONVERT_CHECKPOINT_PATH/model.py \
    --load_platform=huggingface \
    --save_platform=mcore \
    --common_config_path=$CONVERT_CHECKPOINT_PATH/config/llava-ov-1.5-4b/vision-model.json \
    --tensor_model_parallel_size=$TP \
    --load_ckpt_path=$LOAD \
    --save_ckpt_path=$SAVE_VISION_MODEL

# 步骤3：转换适配器
python $CONVERT_CHECKPOINT_PATH/custom/llavaov_1_5/adapter.py \
    --load_platform=huggingface \
    --save_platform=mcore \
    --common_config_path=$CONVERT_CHECKPOINT_PATH/config/llava-ov-1.5-4b/adapter.json \
    --tensor_model_parallel_size=$TP \
    --load_ckpt_path=$LOAD \
    --save_ckpt_path=$SAVE_ADAPTER

# 步骤4：转换视觉补丁嵌入
python $CONVERT_CHECKPOINT_PATH/custom/llavaov_1_5/vision_patch.py \
    --load_platform=huggingface \
    --save_platform=mcore \
    --common_config_path=$CONVERT_CHECKPOINT_PATH/config/llava-ov-1.5-4b/vision-patch.json \
    --load_ckpt_path=$LOAD \
    --save_ckpt_path=$SAVE_PATCH

# 步骤5：合并所有组件
python $CONVERT_CHECKPOINT_PATH/custom/llavaov_1_5/merge_megatron.py \
    --megatron_path $AIAK_MAGATRON_PATH \
    --language_model_path $SAVE_LANGUAGE_MODEL/release \
    --vision_model_path $SAVE_VISION_MODEL/release \
    --vision_patch $SAVE_PATCH/release \
    --adapter_path $SAVE_ADAPTER/release \
    --save_ckpt_path $SAVE/release \
    --tensor_model_parallel_size $TP \
    --pipeline_model_parallel_size $PP

#!/bin/bash
LOAD=$1    # Megatron模型路径
SAVE=$2    # HuggingFace输出路径
TP=$3      # 源张量并行度
PP=$4      # 源流水线并行度

# 步骤1：转换语言模型
python $CONVERT_CHECKPOINT_PATH/model.py \
    --load_platform=mcore \
    --save_platform=huggingface \
    --common_config_path=$CONVERT_CHECKPOINT_PATH/config/llava-ov-1.5-4b/qwen3.json \
    --tensor_model_parallel_size=$TP \
    --pipeline_model_parallel_size=$PP \
    --load_ckpt_path=$LOAD \
    --save_ckpt_path=$SAVE_LANGUAGE_MODEL

# 步骤2-4：转换其他组件（类似）
# ...

# 步骤5：合并到HuggingFace格式
python $CONVERT_CHECKPOINT_PATH/custom/llavaov_1_5/merge_huggingface.py \
    --language_model_path $SAVE_LANGUAGE_MODEL \
    --vision_model_path $SAVE_VISION_MODEL \
    --adapter_path $SAVE_ADAPTER \
    --vision_patch $SAVE_PATCH \
    --save_ckpt_path $SAVE

# 复制非模型文件（tokenizer等）
find $ORIGINAL_HF_PATH/ -type f -not -iname '*safetensors*' -exec cp {} $SAVE/ ';'

{
  "model_type": "GPT",
  "num_layers": 36,
  "hidden_size": 2560,
  "ffn_hidden_size": 9728,
  "num_attention_heads": 32,
  "num_query_groups": 8,
  
  "weight_mapping": {
    "embedding.word_embeddings.weight": "model.embed_tokens.weight",
    "decoder.layers.{}.self_attention.linear_qkv.weight": [
      "model.layers.{}.self_attn.q_proj.weight",
      "model.layers.{}.self_attn.k_proj.weight", 
      "model.layers.{}.self_attn.v_proj.weight"
    ],
    "decoder.layers.{}.self_attention.linear_proj.weight": 
      "model.layers.{}.self_attn.o_proj.weight",
    "decoder.layers.{}.mlp.linear_fc1.weight": [
      "model.layers.{}.mlp.gate_proj.weight",
      "model.layers.{}.mlp.up_proj.weight"
    ],
    "decoder.layers.{}.mlp.linear_fc2.weight": 
      "model.layers.{}.mlp.down_proj.weight",
    "output_layer.weight": "lm_head.weight"
  }
}

{
  "model_type": "VisionModel",
  "num_layers": 24,
  "hidden_size": 1024,
  "num_attention_heads": 16,
  
  "weight_mapping": {
    "decoder.layers.{}.self_attention.linear_qkv.weight": [
      "visual.blocks.{}.attn.qkv.weight"
    ],
    "decoder.layers.{}.self_attention.linear_proj.weight":
      "visual.blocks.{}.attn.proj.weight",
    "decoder.layers.{}.pre_mlp_layernorm.weight":
      "visual.blocks.{}.norm2.weight",
    "decoder.layers.{}.mlp.linear_fc1.weight":
      "visual.blocks.{}.mlp.fc1.weight",
    "decoder.layers.{}.mlp.linear_fc2.weight":
      "visual.blocks.{}.mlp.fc2.weight"
  }
}

def convert_checkpoint(load_platform, save_platform, config, load_path, save_path):
    """执行检查点转换"""
    
    # 1. 加载源检查点
    if load_platform == "huggingface":
        src_checkpoint = load_huggingface_checkpoint(load_path)
    elif load_platform == "mcore":
        src_checkpoint = load_mcore_checkpoint(load_path, config)
    
    # 2. 创建目标检查点结构
    dst_checkpoint = {}
    
    # 3. 根据映射规则转换权重
    weight_mapping = config["weight_mapping"]
    
    for dst_key_template, src_key_template in weight_mapping.items():
        # 处理层索引
        if "{}" in dst_key_template:
            for layer_idx in range(config["num_layers"]):
                dst_key = dst_key_template.format(layer_idx)
                
                if isinstance(src_key_template, list):
                    # 需要拼接多个源权重
                    src_keys = [k.format(layer_idx) for k in src_key_template]
                    weights = [src_checkpoint[k] for k in src_keys]
                    dst_checkpoint[dst_key] = torch.cat(weights, dim=0)
                else:
                    src_key = src_key_template.format(layer_idx)
                    dst_checkpoint[dst_key] = src_checkpoint[src_key]
        else:
            # 直接映射
            dst_checkpoint[dst_key_template] = src_checkpoint[src_key_template]
    
    # 4. 处理张量并行切分（如果需要）
    if save_platform == "mcore" and config.get("tensor_model_parallel_size", 1) > 1:
        dst_checkpoint = split_for_tensor_parallel(
            dst_checkpoint,
            config["tensor_model_parallel_size"]
        )
    
    # 5. 保存目标检查点
    if save_platform == "huggingface":
        save_huggingface_checkpoint(dst_checkpoint, save_path)
    elif save_platform == "mcore":
        save_mcore_checkpoint(dst_checkpoint, save_path, config)

TP=1
PP=1
MBS=1
GBS=8
SEQ_LEN=32768
RECOMPUTE_LAYERS=4

# 内存占用：约60GB/GPU
# 训练速度：约1.2 samples/s

TP=2
PP=1
MBS=1
GBS=8
SEQ_LEN=32768
RECOMPUTE_LAYERS=8

# 内存占用：约70GB/GPU
# 训练速度：约0.8 samples/s

TP=4
PP=2
MBS=1
GBS=16
SEQ_LEN=32768
RECOMPUTE_LAYERS=16

# 内存占用：约75GB/GPU
# 训练速度：约0.3 samples/s

--attention-backend flash

--bf16

--recompute-granularity full
--recompute-method uniform
--recompute-num-layers <根据模型调整>

--use-distributed-optimizer

--num-workers 16

--micro-batch-size 1  # 最小值

--recompute-num-layers 8  # 或更多

--tensor-model-parallel-size 2  # 或更大

--sequence-parallel

--attention-backend flash

--num-workers 16          # 增加工作线程
--dataloader-type cyclic  # 使用循环数据加载器

--bf16  # 或 --fp16

# 平衡TP和DP
# TP太大会增加通信开销
# DP太大可能导致批次大小过小

--lr 5.0e-5  # 从更小的值开始

--clip-grad 1.0

# 确保数据中没有异常值

--optimizer adam
--adam-beta1 0.9
--adam-beta2 0.999
--adam-eps 1e-08

LOGGING_ARGS=(
    --tensorboard-dir ./tensorboard
    --log-timers-to-tensorboard
    --log-memory-to-tensorboard
)

export WANDB_API_KEY=<your-api-key>

LOGGING_ARGS=(
    --wandb-project llava-onevision-1.5
    --wandb-exp-name stage1-alignment
)

LOGGING_ARGS=(
    --log-interval 1              # 日志记录间隔
    --detail-log-interval 20      # 详细日志间隔
    --timing-log-level 0          # 计时日志级别
)