def patch_if_not_exist(group_or_parser, keyname, type=None, default=None):
    """
    只在参数不存在时添加，避免与Megatron冲突
    """
    has_keyname = False
    for action in vars(group_or_parser)["_actions"]:
        if keyname in action.option_strings:
            has_keyname = True
    if not has_keyname:
        return group_or_parser.add_argument(keyname, type=type, default=default)

def initialize_megatron(extra_args_provider=None, args_defaults={}):
    """
    初始化Megatron并验证patch参数
    """
    # 1. 解析参数
    args = parse_args(extra_args_provider, ignore_unknown_args)
    
    # 2. 验证原生Megatron参数
    validate_args(args, args_defaults)
    
    # 3. 验证patch新增的MoE参数
    validate_moe_args(args, args_defaults)
    
    # 4. 设置全局变量
    set_global_variables(args)
    
    # 5. 初始化分布式通信
    _initialize_distributed()

def pretrain(
    train_valid_test_dataset_provider,  # 数据提供函数
    model_provider,                     # 模型构建函数
    model_type,                         # 模型类型枚举
    forward_step_func,                  # 前向传播函数
    process_non_loss_data_func=None,    # 非loss数据处理
    extra_args_provider=None,           # 额外参数注入
    args_defaults={}                    # 参数默认值
):

def finetune(
    train_valid_datasets_provider,     # 数据集提供
        model_provider,                    # 模型构建
    model_type=ModelType.encoder_or_decoder,
    forward_step=_cross_entropy_forward_step,  # 默认CE loss
    end_of_epoch_callback_provider=None,  # epoch结束回调
    task_collate_fn=None               # 自定义collate
):

def model_provider(pre_process=True, post_process=True):
    """
    模型构建函数 - 在每个rank上调用
    """
    args = get_args()
    
    # 1. 构建tokenizer
    build_tokenizer(args)
    
    # 2. 从args构建transformer config
    config = core_transformer_config_from_args(args)
    
    # 3. 获取layer spec（定义层结构）
    transformer_layer_spec = get_gpt_layer_with_transformer_engine_spec(
        args.num_experts, args.moe_grouped_gemm
    )
    
    # 4. 构建模型
    model = GPTModel(
        config=config,
        transformer_layer_spec=transformer_layer_spec,
        vocab_size=args.padded_vocab_size,
        max_sequence_length=args.max_position_embeddings,
        pre_process=pre_process,     # 首个PP stage包含embedding
        post_process=post_process,   # 最后PP stage包含output layer
        fp16_lm_cross_entropy=args.fp16_lm_cross_entropy,
        parallel_output=True,
        share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights,
        position_embedding_type=args.position_embedding_type,
        rotary_percent=args.rotary_percent,
        rotary_base=args.rotary_base,
    )
    
    return model

def get_gpt_layer_with_transformer_engine_spec(num_experts=None):
    """
    使用Transformer Engine加速的layer spec
    """
    return ModuleSpec(
        module=TransformerLayer,
        submodules=TransformerLayerSubmodules(
            self_attention=ModuleSpec(
                module=SelfAttention,
                params={"attn_mask_type": AttnMaskType.causal},
                submodules=SelfAttentionSubmodules(
                    linear_qkv=TELayerNormColumnParallelLinear,  # TE融合算子
                    core_attention=TEDotProductAttention,        # TE FlashAttention
                    linear_proj=TERowParallelLinear,
                    q_layernorm=TENorm if qk_layernorm else IdentityOp,
                    k_layernorm=TENorm if qk_layernorm else IdentityOp,
                ),
            ),
            self_attn_bda=get_bias_dropout_add,  # 融合bias+dropout+add
            pre_mlp_layernorm=TENorm,
            mlp=mlp,
            mlp_bda=get_bias_dropout_add,
        ),
    )

class MultiLatentAttention(Attention):
    def __init__(self, config: MLATransformerConfig, ...):
        # Q投影：分为两步
        self.linear_q_down_proj: hidden → q_lora_rank  # 降维
        self.linear_q_up_proj: q_lora_rank → num_heads * qk_head_dim  # 升维
        
        # KV投影：共享latent space
        self.linear_kv_down_proj: hidden → kv_lora_rank  # 压缩
        self.linear_kv_up_proj: kv_lora_rank → num_heads * (qk + v)_head_dim  # 解压
        
        # RoPE只应用于qk_pos_emb_head_dim维度
        self.rotary_pos_emb = YarnRotaryEmbedding(
            qk_pos_emb_head_dim,
            rotary_base=config.rotary_base,
            scaling_factor=config.rotary_scaling_factor,
            ...
        )

# megatron_patch/model/qwen2_moe/layer_specs.py
def get_gpt_layer_with_transformer_engine_spec_moe():
    mlp = ModuleSpec(
        module=MoELayer,
        submodules=MoESubmodules(
            experts=ModuleSpec(
                module=TEGroupedMLP,  # Grouped GEMM优化
                submodules=MLPSubmodules(
                    linear_fc1=TEColumnParallelLinear,
                    linear_fc2=TERowParallelLinear,
                ),
            ),
            router=TopKRouter,  # Top-K路由
            shared_experts=ModuleSpec(...),  # Qwen2-MoE的shared experts
        )
    )

# 工具: toolkits/pretrain_data_preprocessing/preprocess_data_megatron.py
python preprocess_data_megatron.py \
    --input my_corpus.json \
    --output-prefix my_corpus_tokenized \
    --tokenizer-type GPT2BPETokenizer \
    --vocab-file vocab.json \
    --merge-file merges.txt \
    --append-eod \
    --workers 32

def train_valid_test_datasets_provider(train_val_test_num_samples):
    args = get_args()
    
    if "-Raw" in args.dataset:
        # 原始格式：每次训练时tokenize
        train_ds, valid_ds, test_ds = build_pretrain_dataset_from_original(args.dataset)
    elif "-Idxmap" in args.dataset:
        # Megatron索引格式：预先tokenize
        config = GPTDatasetConfig(
            random_seed=args.seed,
            sequence_length=args.seq_length,
            blend=args.data_path,  # 支持blend多数据集
            split=args.split,  # train:valid:test比例
            eod_mask_loss=args.eod_mask_loss,
        )
        train_ds, valid_ds, test_ds = BlendedMegatronDatasetBuilder(
            GPTDataset, train_val_test_num_samples, config
        ).build()
    
    return train_ds, valid_ds, test_ds

{
  "id": "example_001",
  "conversations": [
    {"from": "human", "value": "你好，请介绍一下你自己"},
    {"from": "gpt", "value": "我是一个AI助手，可以帮助你..."}
  ]
}

class TaskEncoder(DefaultTaskEncoder):
    def encode_chatml(self, sample: ChatMLSample):
        # 1. 图像预处理
        imgs = [get_visual_transform(img) for img in sample.imgs]
        flattened_imgs, image_thw_grids = self._flatten_visual_inputs(imgs)
        
        # 2. 构建conversation
        conversation = [
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": [
                {"type": "image", "image": "0"},
                {"type": "text", "text": "描述这张图片"}
            ]},
            {"role": "assistant", "content": "这是一张..."}
        ]
        
        # 3. Tokenize并插入image placeholders
        input_ids = tokenizer.apply_chat_template(conversation, tokenize=True)
        
        # 4. 扩展<|image_pad|> token
        # 一个image → (grid_t * grid_h * grid_w) // merge_size² 个tokens
        for image_token_idx in image_token_indices:
            size = image_thw_grids[i].prod() // merge_length
            # 替换单个<|image_pad|>为size个token
            final_input_ids[cur:cur+size] = image_token_id
        
        return ImageTaskSample(
            imgs=flattened_imgs,
            image_thw_grids=image_thw_grids,
            text=final_input_ids,
            target=labels,
            image_input_mask=image_mask,
        )

def convert_checkpoint_from_transformers_to_megatron(args):
    # 1. 加载HF模型
    hf_model = AutoModelForCausalLM.from_pretrained(args.load_path)
    hf_state_dict = hf_model.state_dict()
    
    # 2. 初始化Megatron模型（指定TP/PP/EP）
    mg_model = model_provider()
    
    # 3. 权重映射和切分
    mg_state_dict = {}
    for hf_name, hf_param in hf_state_dict.items():
        # 命名空间转换
        mg_name = convert_name_hf_to_mg(hf_name)
        
        # 张量切分（TP维度）
        if "self_attn.q_proj" in hf_name:
            # Split along num_heads dimension
            mg_param = split_tensor_for_tp(hf_param, args.target_tensor_model_parallel_size)
        elif "mlp.up_proj" in hf_name:
            # Column parallel: split along hidden_size
            mg_param = split_tensor_for_tp(hf_param, args.target_tensor_model_parallel_size)
        elif "mlp.down_proj" in hf_name:
            # Row parallel: split along input dimension
            mg_param = split_tensor_for_tp(hf_param, args.target_tensor_model_parallel_size, dim=1)
        
        mg_state_dict[mg_name] = mg_param
    
    # 4. 保存为Megatron checkpoint格式
    for tp_rank in range(args.target_tensor_model_parallel_size):
        for pp_rank in range(args.target_pipeline_model_parallel_size):
            save_checkpoint(
                iteration, mg_model, None, None,
                checkpoint_name=get_checkpoint_name(load_dir, iteration, release, tp_rank, pp_rank)
            )

# HuggingFace → Megatron
{
    "model.embed_tokens.weight": "embedding.word_embeddings.weight",
    "model.layers.0.self_attn.q_proj.weight": "decoder.layers.0.self_attention.linear_qkv.weight[0:hidden]",
    "model.layers.0.self_attn.k_proj.weight": "decoder.layers.0.self_attention.linear_qkv.weight[hidden:2*hidden]",
    "model.layers.0.self_attn.v_proj.weight": "decoder.layers.0.self_attention.linear_qkv.weight[2*hidden:3*hidden]",
    "model.layers.0.self_attn.o_proj.weight": "decoder.layers.0.self_attention.linear_proj.weight",
    "model.layers.0.mlp.gate_proj.weight": "decoder.layers.0.mlp.linear_fc1.weight[0:intermediate]",
    "model.layers.0.mlp.up_proj.weight": "decoder.layers.0.mlp.linear_fc1.weight[intermediate:2*intermediate]",
    "model.layers.0.mlp.down_proj.weight": "decoder.layers.0.mlp.linear_fc2.weight",
}

# impl/convert.py
class DistributedCheckpointConverter:
    def convert(self, src_path, dst_path, src_type, dst_type):
        """
        并行转换checkpoint
        
        支持格式：
        - HuggingFace safetensors
        - Megatron legacy (tp-pp-ep组织)
        - torch_dist (PyTorch Distributed format)
        """
        # 1. 构建权重映射表（不加载权重）
        weight_map = self.build_weight_mapping(src_type, dst_type)
        
        # 2. 分片加载和转换
        for shard_id in parallel_range(num_shards):
            # 只加载当前shard需要的权重
            shard_weights = load_shard(src_path, shard_id, weight_map)
            
            # 转换并写入目标shard
            converted = self.convert_shard(shard_weights, weight_map)
            save_shard(dst_path, shard_id, converted)

# 64个experts，EP=8，每个rank持有8个experts
def split_experts_for_ep(expert_weights, num_experts, ep_size):
    experts_per_rank = num_experts // ep_size
    for ep_rank in range(ep_size):
        start_expert = ep_rank * experts_per_rank
        end_expert = (ep_rank + 1) * experts_per_rank
        
        rank_experts = {}
        for expert_id in range(start_expert, end_expert):
            rank_experts[f"expert_{expert_id}"] = expert_weights[expert_id]
        
        save_checkpoint_for_ep_rank(ep_rank, rank_experts)

# examples/qwen3/run_mcore_qwen.sh

# 分布式配置
GPUS_PER_NODE=8
MASTER_ADDR=localhost
MASTER_PORT=6000
NNODES=4
NODE_RANK=0
WORLD_SIZE=$(($GPUS_PER_NODE*$NNODES))

# 并行策略
TP=4    # Tensor Parallel
PP=2    # Pipeline Parallel
EP=1    # Expert Parallel（MoE模型）
CP=1    # Context Parallel

# 模型配置
SEQ_LENGTH=4096
HIDDEN_SIZE=4096
NUM_LAYERS=32
NUM_ATTENTION_HEADS=32

# 训练超参数
GLOBAL_BATCH_SIZE=256
MICRO_BATCH_SIZE=1
GRADIENT_ACCUMULATION_STEPS=$((GLOBAL_BATCH_SIZE / (MICRO_BATCH_SIZE * WORLD_SIZE / TP / PP)))

# 启动命令
torchrun --nproc_per_node $GPUS_PER_NODE \
         --nnodes $NNODES \
         --node_rank $NODE_RANK \
         --master_addr $MASTER_ADDR \
         --master_port $MASTER_PORT \
    pretrain_qwen.py \
        --tensor-model-parallel-size $TP \
        --pipeline-model-parallel-size $PP \
        --expert-model-parallel-size $EP \
        --context-parallel-size $CP \
        --num-layers $NUM_LAYERS \
        --hidden-size $HIDDEN_SIZE \
        --num-attention-heads $NUM_ATTENTION_HEADS \
        --seq-length $SEQ_LENGTH \
        --max-position-embeddings $SEQ_LENGTH \
        --micro-batch-size $MICRO_BATCH_SIZE \
        --global-batch-size $GLOBAL_BATCH_SIZE \
        --train-iters 10000 \
        --lr 1e-4 \
        --min-lr 1e-5 \
        --lr-decay-style cosine \
        --weight-decay 0.1 \
        --adam-beta1 0.9 \
        --adam-beta2 0.95 \
        --clip-grad 1.0 \
        --init-method-std 0.01 \
        --attention-dropout 0.0 \
        --hidden-dropout 0.0 \
        --no-query-key-layer-scaling \
        --position-embedding-type rope \
        --rotary-base 1000000 \
        --rotary-percent 1.0 \
        --use-rotary-position-embeddings \
        --swiglu \
        --normalization RMSNorm \
        --untie-embeddings-and-output-weights \
        --disable-bias-linear \
        --use-flash-attn \
        --data-path /path/to/data_text_document \
        --split 949,50,1 \
        --tokenizer-type PretrainedFromHF \
        --tokenizer-name-or-path /path/to/tokenizer \
        --save-interval 500 \
        --eval-interval 100 \
        --log-interval 1 \
        --save /path/to/checkpoints \
        --load /path/to/checkpoints \
        --tensorboard-dir /path/to/tensorboard \
        --no-load-optim \
        --no-load-rng \
        --bf16

def forward_step(data_iterator, model):
    """
    单个microbatch的前向传播
    """
    timers = get_timers()
    
    # 1. 获取数据
    timers('batch-generator').start()
    tokens, labels, loss_mask, attention_mask, position_ids, _ = get_batch(data_iterator)
    timers('batch-generator').stop()
    
    # 2. 模型前向
    output_tensor = model(
        input_ids=tokens,
        position_ids=position_ids,
        attention_mask=attention_mask,
        labels=labels
    )
    
    # 3. 计算loss（只在最后一个PP stage）
    if mpu.is_pipeline_last_stage():
        loss = compute_loss(output_tensor, labels, loss_mask)
        return loss, {'lm loss': loss}
    else:
        return output_tensor

# Megatron使用的DistributedOptimizer
optimizer = get_megatron_optimizer(model)

# 特点：
# 1. 参数和梯度自动处理TP/PP分布
# 2. 梯度all-reduce跨DP组
# 3. 支持optimizer state的DP分片（ZeRO-1）

def get_learning_rate_scheduler(optimizer):
    """
    常用：Cosine Decay with Warmup
    """
    if args.lr_decay_style == 'cosine':
        lr_scheduler = CosineAnnealingScheduler(
            optimizer,
            max_lr=args.lr,
            min_lr=args.min_lr,
            warmup_steps=args.lr_warmup_iters,
            decay_steps=args.lr_decay_iters,
        )
    return lr_scheduler

# 保存
save_checkpoint(iteration, model, optimizer, opt_param_scheduler)
# 生成：/save/iter_0001000/mp_rank_XX_YY/model_optim_rng.pt
#     XX=tp*pp+pp_rank, YY=000

# 加载
load_checkpoint(model, optimizer, opt_param_scheduler)
# 自动从latest_checkpointed_iteration.txt读取iteration
# 根据当前rank的tp/pp位置加载对应shard

# examples/qwen3/chatlearn_entrypoint.py
class ChatlearnLauncher:
    def launch(self):
        # 定义4个模型的角色
        models = {
            "policy": PolicyModel(),      # 待训练的Actor
            "reference": ReferenceModel(), # 冻结的参考模型
            "reward": RewardModel(),       # 打分模型
            "value": ValueModel(),         # 价值估计（PPO用）
        }
        
        # 定义计算流程
        def rlhf_workflow(env, policy, reference, reward, value):
            # 1. Rollout：生成responses
            prompts = env.get_prompts()
            responses = policy.generate(prompts, num_samples=4)  # 每个prompt生成4个
            
            # 2. Reward：计算分数
            scores = reward.score(prompts, responses)
            
            # 3. Reference log_probs：KL正则化
            ref_log_probs = reference.forward(prompts, responses)
            
            # 4. 计算advantages（GRPO相对优势）
            group_mean_score = scores.reshape(-1, 4).mean(dim=1)
            advantages = scores - group_mean_score.repeat_interleave(4)
            
            # 5. Policy更新
            policy_loss = compute_policy_loss(
                policy.forward(prompts, responses),
                ref_log_probs,
                advantages
            )
            policy.backward(policy_loss)
        
        # 启动ChatLearn引擎
        engine = ChatLearnEngine(models, rlhf_workflow)
        engine.train(num_episodes=1000)

# examples/qwen3/verl_entrypoint.py
@ray.remote
class TaskRunner:
    def run(self, config):
        # 1. 初始化模型workers
        actor_rollout_worker = RayWorkerGroup(
            cls=MegatronPPOActor,
            config=config.actor_rollout_ref,
        )
        
        critic_worker = RayWorkerGroup(
            cls=MegatronCritic,
            config=config.critic,
        )
        
        reward_model = load_reward_model(config.reward_model)
        
        # 2. 创建trainer
        trainer = RayPPOTrainer(
            config=config,
            actor_rollout_worker=actor_rollout_worker,
            critic_worker=critic_worker,
            reward_model=reward_model,
        )
        
        trainer.fit()

# verl_patch/trainer/ray_trainer.py
class RayPPOTrainer:
    def fit(self):
        for epoch in range(total_epochs):
            for batch in train_dataloader:
                # 1. Rollout阶段
                with torch.no_grad():
                    rollout_output = self.actor_rollout.generate_sequences(batch)
                    # rollout_output: {sequences, log_probs, values}
                
                # 2. Reward计算
                rewards = self.reward_model.compute_rewards(
                    rollout_output['sequences']
                )
                
                # 3. Advantage计算（GAE）
                advantages, returns = compute_gae(
                    rewards,
                    rollout_output['values'],
                    gamma=0.99,
                    lambda_=0.95
                )
                
                # 4. PPO更新（多个epochs）
                for ppo_epoch in range(self.config.ppo_epochs):
                    # Actor更新
                    actor_loss = self.compute_actor_loss(
                        rollout_output,
                        advantages,
                        clip_ratio=0.2
                    )
                    self.actor.backward(actor_loss)
                    
                    # Critic更新
                    critic_loss = self.compute_critic_loss(
                        rollout_output['sequences'],
                        returns
                    )
                    self.critic.backward(critic_loss)

--use-flash-attn  # 启用FlashAttention-2
--use-te-flash-attn  # 使用Transformer Engine的FlashAttention实现

submodules=SelfAttentionSubmodules(
    core_attention=TEDotProductAttention,  # TE backend自动使用FA
    # or
    core_attention=DotProductAttention,    # Megatron实现，支持--use-flash-attn
)

--transformer-impl transformer_engine \
--fp8-format hybrid  # E4M3 (forward) + E5M2 (backward) \
--fp8-amax-history-len 1024 \
--fp8-amax-compute-algo max

# 需要cu_seqlens指示每个样本的边界
cu_seqlens = torch.tensor([0, len1, len1+len2, len1+len2+len3])

# FlashAttention支持variable-length sequences
attn_output = flash_attn_varlen_func(
    q, k, v,
    cu_seqlens_q=cu_seqlens,
    cu_seqlens_k=cu_seqlens,
    max_seqlen_q=max_len,
    max_seqlen_k=max_len,
)

--context-parallel-size 4  # 将seq_len切分到4个GPU
--sequence-parallel  # 必须同时启用SP

# 将sequence维度切分
seq_len_per_rank = seq_len // cp_size

# Ring Attention: 每个rank持有部分Q和完整KV
# 通过ring communication访问其他rank的KV
for cp_rank in range(cp_size):
    local_q = q_full[cp_rank * seq_chunk : (cp_rank+1) * seq_chunk]
    
    for step in range(cp_size):
        # 当前持有的KV shard
        kv_rank = (cp_rank + step) % cp_size
        local_kv = kv_shards[kv_rank]
        
        # 计算local attention
        local_attn = torch.matmul(local_q, local_kv.transpose(-1, -2))
        
        # 异步发送KV到下一个rank，接收前一个rank的KV
        send_recv_kv(local_kv, next_rank, prev_rank)

--use-distributed-optimizer  # 启用分布式优化器
--overlap-grad-reduce  # 梯度reduce与backward overlap
--offload-optimizer  # 优化器状态offload到CPU

# 优化器状态（Adam: momentum + variance）存储在CPU
optimizer_state = {
    "exp_avg": tensor.to("cpu"),      # momentum
    "exp_avg_sq": tensor.to("cpu"),   # variance
}

# Backward时梯度累积在GPU
grads_gpu = accumulate_gradients()

# Optimizer step前将需要的state传回GPU
for param in model.parameters():
    state = optimizer_state[param].to("gpu")
    param.data = adam_update(param.data, grad, state)
    optimizer_state[param] = state.to("cpu")

python report_auto_config.py \
    --model-size 70B \
    --num-gpus 64 \
    --gpu-memory 80GB \
    --seq-length 4096 \
    --global-batch-size 256

# 清洗
python clean_raw_text.py \
    --input raw_corpus.txt \
    --output cleaned_corpus.txt \
    --remove-duplicate-lines \
    --min-length 50

# Tokenize
python preprocess_data_megatron.py \
    --input cleaned_corpus.txt \
    --output-prefix my_data \
    --tokenizer-type PretrainedFromHF \
    --tokenizer-name-or-path Qwen/Qwen-7B \
    --workers 64 \
    --append-eod

# 输出: my_data_text_document.{bin,idx}

python build_idxmap_sft_dataset.py \
    --input sft_data.jsonl \
    --output-prefix sft_indexed \
    --tokenizer-name-or-path Qwen/Qwen-7B \
    --seq-length 2048 \
    --apply-chat-template \
    --mask-instruction

# LLaVA format → WebDataset
python convert_llava_pretrain_to_wds.py \
    --input llava_pretrain.json \
    --image-folder /path/to/images \
    --output wds_shards/shard-{000000..000099}.tar \
    --max-shard-size 1GB

python evaluate_megatron_model.py \
    --model-path /path/to/checkpoint \
    --tasks mmlu,ceval,humaneval \
    --num-fewshot 5 \
    --tensor-model-parallel-size 4 \
    --pipeline-model-parallel-size 2 \
    --output-path eval_results.json

# 1. Gradient Explosion
--clip-grad 1.0  # 梯度裁剪
--init-method-std 0.01  # 更小的初始化标准差

# 2. Loss Spike
--lr-warmup-iters 2000  # 充足的warmup
--lr-decay-style cosine  # 平滑的LR decay
--attention-dropout 0.1  # 适度dropout
--hidden-dropout 0.1

# 3. NaN/Inf
--check-for-nan-in-loss-and-grad  # 自动检测
--bf16  # BF16比FP16更稳定（动态范围大）
--initial-loss-scale 4096  # FP16时调整loss scale

# 4. MoE Load Imbalance
--moe-aux-loss-coeff 0.01  # Auxiliary loss权重
--moe-z-loss-coeff 0.001   # Z-loss稳定router

# 频繁保存（防止训练中断）
--save-interval 500

# 但不要每次都保存optimizer state（节省空间）
--no-save-optim  # 每隔N次保存一次完整checkpoint

# 加载时可选择性恢复
--load /path/to/checkpoint
--no-load-optim  # 从pretrain checkpoint开始SFT时
--no-load-rng    # 改变batch size时

# Level 0: Baseline
micro_batch_size = 4

# Level 1: 激活检查点（Activation Checkpointing）
--recompute-activations  # 减少50% activation memory, +20% time

# Level 2: Sequence Parallel
--sequence-parallel  # 与TP配合，分摊LayerNorm/Dropout activation

# Level 3: 更小的micro batch
micro_batch_size = 1  # 配合更多gradient accumulation

# Level 4: CPU Offloading
--use-distributed-optimizer
--offload-optimizer  # 减少30-50% memory, +10-15% time

# Level 5: 混合策略
--recompute-granularity selective  # 只checkpoint部分层
--zero-stage 1  # ZeRO stage 1: optimizer state sharding

# 1. Flash Attention (必开)
--use-flash-attn

# 2. Transformer Engine (Ampere/Hopper)
--transformer-impl transformer_engine

# 3. Overlap通信与计算
--overlap-grad-reduce  # Backward时overlap
--overlap-param-gather  # ZeRO时overlap

# 4. 编译优化
--use-torch-compile  # PyTorch 2.0+

# 5. Fused Kernels
--bias-dropout-fusion
--bias-swiglu-fusion  # LLaMA SwiGLU
--apply-layernorm-1p  # Fused LayerNorm

# 6. 数据加载
--num-workers 4  # DataLoader多进程
--dataloader-type cyclic  # 预加载下一个epoch

# 1. 准备HuggingFace模型权重
# (如果从头预训练，可跳过此步)

# 2. 转换为Megatron格式
bash toolkits/model_checkpoints_convertor/your_model/hf2mcore_convertor.sh

# 3. 准备训练数据
# 预训练
bash toolkits/pretrain_data_preprocessing/run_make_pretraining_dataset_megatron.sh

# SFT
bash toolkits/sft_data_preprocessing/run_build_idxmap_sft_dataset.sh

# 4. 自动配置
python toolkits/auto_configurator/report_auto_config.py \
    --model-size 7B --num-gpus 8 --gpu-memory 80GB

# 5. 修改训练脚本
cp examples/llama2/pretrain_mcore_llama.py examples/your_model/
cp examples/llama2/run_pretrain_mcore_llama.sh examples/your_model/
# 修改模型配置和并行参数

# 6. 启动训练
bash examples/your_model/run_pretrain_mcore_llama.sh

# 7. 监控训练
tensorboard --logdir /path/to/tensorboard_dir

# 8. 评估
bash examples/your_model/run_evaluate_mcore.sh

# 9. 转换回HuggingFace格式
bash toolkits/model_checkpoints_convertor/your_model/mcore2hf_convertor.sh

# 10. 部署推理
python examples/your_model/generate_text.py

# 检查项：
1. micro_batch_size是否过大？
   → 减小到1，增加gradient accumulation

2. 是否启用了activation checkpointing？
   → --recompute-activations

3. TP/PP是否合理？
   → 增大TP以分摊模型参数

4. Sequence length是否过长？
   → 考虑--context-parallel-size

5. 是否有memory leak？
   → torch.cuda.empty_cache() 定期清理

# 检查项：
1. Learning rate是否合理？
   → 典型值：7B用1e-4, 70B用6e-5
   → Warmup是否充足？建议≥2000 steps

2. 数据质量？
   → 检查是否有重复/错误样本
   → 验证tokenization正确性

3. 梯度是否异常？
   → --log-grad-norm 查看梯度范数
   → 是否出现NaN/Inf？

4. 模型配置是否正确？
   → 检查vocab_size, hidden_size等是否匹配HF config

# 排查项：
1. GPU利用率低？
   → nvidia-smi查看GPU util
   → 如果<80%，检查数据加载是否成为瓶颈

2. 通信时间长？
   → nsys profile查看通信占比
   → TP过大或跨节点TP都会导致通信瓶颈

3. 是否启用了所有优化？
   → FlashAttention, Transformer Engine, Overlap等

4. Micro batch size过小？
   → 增大micro_batch以提升GPU利用率

# 启动分布式训练（单节点）
torchrun --nproc_per_node 8 pretrain.py [args]

# 启动分布式训练（多节点）
torchrun --nnodes 4 --node_rank $RANK \
         --master_addr $MASTER_ADDR \
         --master_port $MASTER_PORT \
         --nproc_per_node 8 \
         pretrain.py [args]

# 权重转换 HF→Megatron
bash hf2mcore_convertor.sh

# 权重转换 Megatron→HF
bash mcore2hf_convertor.sh

# 数据预处理
python preprocess_data_megatron.py \
    --input data.txt \
    --output-prefix data_tokenized \
    --tokenizer-type PretrainedFromHF \
    --tokenizer-name-or-path /path/to/tokenizer

# 评估
bash run_evaluate_mcore.sh

# 推理生成
bash run_text_generation.sh