1
def load_model(args):
2
    local_rank = int(os.environ.get("LOCAL_RANK", 0))
3
    model, tokenizer = FastLanguageModel.from_pretrained(
4
        model_name=args.model,  # "unsloth/Qwen3.5-35B-A3B"
5
        max_seq_length=args.max_seq_length,  # 32768
6
        load_in_4bit=False,  # MoE + BnB 4bit 不兼容，必须用 bf16
7
        dtype=torch.bfloat16,
8
        device_map={"": local_rank},  # FSDP 需要每个 rank 只加载到自己的 GPU
9
    )
10

11
    # FSDP 要求所有参数 dtype 一致，把 fp32 的参数强制转 bf16
12
    for name, param in model.named_parameters():
13
        if param.dtype == torch.float32:
14
            param.data = param.data.to(torch.bfloat16)
15

16
    # Qwen3.5 是 VLM，tokenizer 可能是 Processor 包装的
17
    tokenizer = get_chat_template(tokenizer, chat_template="qwen-2.5")
18
    text_tokenizer = getattr(tokenizer, "tokenizer", tokenizer)
19

20
    model = FastLanguageModel.get_peft_model(
21
        model,
22
        r=32,
23
        lora_alpha=32,
24
        target_modules=[
25
            "q_proj", "k_proj", "v_proj", "o_proj",
26
            "gate_proj", "up_proj", "down_proj",
27
        ],
28
        lora_dropout=0,
29
        bias="none",
30
        use_gradient_checkpointing="unsloth",
31
        random_state=3407,
32
    )
33
    return model, tokenizer, text_tokenizer

1
def _patch_model_with_cce(model):
2
    """替换 lm_head 为 identity，用 CCE 直接算 loss"""
3
    import types
4
    from transformers.modeling_outputs import CausalLMOutputWithPast
5

6
    # 层层剥开 PEFT 包装，找到真正的 CausalLM 模型
7
    lm_model = model
8
    while hasattr(lm_model, "model"):
9
        lm_model = lm_model.model
10
    if not hasattr(lm_model, "lm_head"):
11
        lm_model = model.base_model.model
12
        if hasattr(lm_model, "language_model"):
13
            lm_model = lm_model.language_model
14

15
    # 保存 lm_head 权重，替换为 identity
16
    lm_head_weight = lm_model.lm_head.weight  # (vocab_size, hidden_dim)
17

18
    class IdentityLMHead(torch.nn.Module):
19
        def __init__(self, weight):
20
            super().__init__()
21
            self.weight = weight
22
        def forward(self, hidden_states):
23
            return hidden_states  # 直接返回 hidden states，不做 vocab projection
24

25
    lm_model.lm_head = IdentityLMHead(lm_head_weight)
26

27
    # Monkey-patch forward
28
    orig_forward = lm_model.forward.__func__
29

30
    def cce_forward(self, *args, **kwargs):
31
        labels = kwargs.pop("labels", None)
32
        kwargs["labels"] = None  # 不让原始 forward 计算 loss
33

34
        outputs = orig_forward(self, *args, **kwargs)
35

36
        if labels is not None:
37
            hidden_states = outputs.logits  # 其实是 identity 返回的 hidden states
38

39
            if hidden_states.dtype == torch.float32:
40
                hidden_states = hidden_states.to(torch.bfloat16)
41
            lm_weight = self.lm_head.weight
42
            if lm_weight.dtype == torch.float32:
43
                lm_weight = lm_weight.to(torch.bfloat16)
44

45
            # 直接从 hidden × weight 计算 CE loss，不分配 logits
46
            loss = linear_cross_entropy(
47
                hidden_states, lm_weight, labels,
48
                ignore_index=-100, shift=True, reduction="mean",
49
            )
50

51
            return CausalLMOutputWithPast(
52
                loss=loss,
53
                logits=None,  # 没有 logits — 省了 ~30GB
54
                past_key_values=outputs.past_key_values,
55
                hidden_states=outputs.hidden_states,
56
                attentions=outputs.attentions,
57
            )
58
        return outputs
59

60
    lm_model.forward = types.MethodType(cce_forward, lm_model)

1
标准流程：
2
  hidden_states → lm_head(Linear) → logits [30GB tensor] → CrossEntropy → loss
3

4
CCE 流程：
5
  hidden_states → identity (pass through) → linear_cross_entropy(hidden, weight, labels) → loss
6
                                              ↑ 无需分配 logits，直接在 hidden 上算

1
def prepare_dataset(args, tokenizer, text_tokenizer):
2
    dataset = load_dataset(args.dataset, split="train")
3

4
    # 过滤：必须有 user 和 assistant 角色
5
    def has_user_role(example):
6
        msgs = json.loads(example["messages"]) if isinstance(example["messages"], str) else example["messages"]
7
        roles = {m["role"] for m in msgs}
8
        return "user" in roles and "assistant" in roles
9

10
    dataset = dataset.filter(has_user_role, num_proc=4)
11

12
    # 子采样
13
    if args.max_samples > 0 and len(dataset) > args.max_samples:
14
        dataset = dataset.shuffle(seed=42).select(range(args.max_samples))
15

16
    # 用 chat template 格式化
17
    def format_sharegpt(example):
18
        messages = json.loads(example["messages"]) if isinstance(example["messages"], str) else example["messages"]
19
        text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)
20
        return {"text": text}
21

22
    dataset = dataset.map(format_sharegpt, num_proc=4, remove_columns=dataset.column_names)
23

24
    # 手动 tokenize — 绕过 SFTTrainer 的 eos_token 验证 bug
25
    def tokenize(examples):
26
        all_ids, all_masks = [], []
27
        for text in examples["text"]:
28
            enc = text_tokenizer(text, truncation=True, max_length=args.max_seq_length)
29
            all_ids.append(enc["input_ids"])
30
            all_masks.append(enc["attention_mask"])
31
        return {"input_ids": all_ids, "attention_mask": all_masks, "labels": all_ids}
32

33
    dataset = dataset.map(tokenize, batched=True, remove_columns=["text"])
34
    return dataset

1
def train(args, model, tokenizer, dataset):
2
    training_args = SFTConfig(
3
        output_dir=args.output_dir,
4
        per_device_train_batch_size=args.batch_size,       # 2
5
        gradient_accumulation_steps=args.grad_accum,       # 4
6
        warmup_ratio=0.03,
7
        num_train_epochs=args.epochs,                      # 1
8
        learning_rate=args.lr,                             # 2e-4
9
        lr_scheduler_type="cosine",
10
        bf16=True,
11
        logging_steps=1,
12
        save_steps=200,
13
        save_total_limit=3,
14
        optim="adamw_torch",
15
        weight_decay=0.01,
16
        max_seq_length=args.max_seq_length,                # 32768
17
        packing=False,
18
        report_to="wandb",
19
        dataloader_num_workers=2,
20
    )
21

22
    # 关键：batch_size > 1 时必须有 DataCollator 来 padding 变长序列
23
    text_tokenizer = getattr(tokenizer, "tokenizer", tokenizer)
24
    data_collator = DataCollatorForSeq2Seq(
25
        tokenizer=text_tokenizer,
26
        padding=True,
27
        pad_to_multiple_of=8,
28
    )
29

30
    trainer = SFTTrainer(
31
        model=model,
32
        tokenizer=tokenizer,
33
        train_dataset=dataset,
34
        data_collator=data_collator,
35
        args=training_args,
36
    )
37
    trainer.train()