Weights & Biases 实验追踪

import wandb

run = wandb.init(
    project="llm-finetuning",       # 项目名称，用于组织实验
    entity="my-team",               # 团队/用户名，默认为个人账号
    config={                        # 超参数配置字典
        "learning_rate": 2e-5,
        "epochs": 3,
        "batch_size": 16,
        "model_name": "meta-llama/Llama-2-7b-hf",
        "lora_rank": 8,
        "lora_alpha": 16,
        "weight_decay": 0.01,
        "warmup_steps": 100,
    },
    name="llama2-lora-exp01",       # 运行名称，在仪表盘中显示
    tags=["lora", "llama2", "sft"], # 标签，方便筛选
    notes="LoRA微调Llama2，使用Alpaca数据集",  # 备注说明
    dir="./wandb_logs",             # 本地日志存储目录
    save_code=True,                 # 保存运行脚本的代码快照
    reinit=True,                    # 允许在同一脚本中多次初始化
    resume="allow",                 # 允许恢复中断的运行（"must"必须恢复，"allow"允许恢复）
    id="unique-run-id-123",         # 自定义运行ID，用于恢复时匹配
    group="lora-ablation",          # 分组名称，用于组织相关实验
    job_type="train",               # 作业类型（train/eval/test）
)

# 访问运行信息
print(f"运行ID: {run.id}")
print(f"运行URL: {run.url}")
print(f"配置: {run.config}")

# 基础用法：记录标量指标
wandb.log({"train/loss": 2.345, "train/accuracy": 0.67})

# 指定步数（step）
wandb.log({"train/loss": 1.876}, step=100)

# 记录多个指标（同一step下）
wandb.log({
    "train/loss": 1.523,
    "train/learning_rate": 1.8e-5,
    "train/perplexity": 4.58,
    "train/tokens_per_sec": 5200,
}, step=200)

# 使用 commit=False 延迟提交（攒一批再提交）
wandb.log({"train/loss": 1.2}, commit=False)
wandb.log({"train/accuracy": 0.85}, commit=True)  # 此时一起提交

# 记录自定义图表
wandb.log({
    "attention_heatmap": wandb.plots.HeatMap(
        x_labels=range(seq_len),
        y_labels=range(seq_len),
        matrix_values=attention_weights,
    )
})

# 训练循环中的典型用法
for epoch in range(num_epochs):
    for step, batch in enumerate(train_dataloader):
        loss = model(batch).loss
        loss.backward()
        optimizer.step()
        scheduler.step()

        if step % 10 == 0:
            wandb.log({
                "train/loss": loss.item(),
                "train/lr": scheduler.get_last_lr()[0],
                "train/epoch": epoch + step / len(train_dataloader),
            }, step=global_step)
            global_step += 1

# 方式1：通过 init 传入
run = wandb.init(config={"lr": 2e-5, "epochs": 3})

# 方式2：通过 argparse + config
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--lr", type=float, default=2e-5)
args = parser.parse_args()
wandb.init(config=vars(args))

# 方式3：运行时动态设置
wandb.config.update({"additional_key": "value"})

# 方式4：从配置文件加载
import yaml
with open("config.yaml") as f:
    config = yaml.safe_load(f)
wandb.init(config=config)

# 访问配置
learning_rate = wandb.config.learning_rate
epochs = wandb.config.epochs

# Sweep 中，config 会被 Sweep Agent 自动更新为超参数组合
# 此时 wandb.config 包含 Sweep 定义的参数值

# ===== 创建并记录模型 Artifact =====
# 方式1：直接记录目录
artifact = wandb.Artifact(
    name="llama2-lora-model",       # Artifact名称
    type="model",                    # 类型：model, dataset, code 等
    description="LoRA微调后的Llama2模型",
    metadata={                       # 元数据
        "base_model": "meta-llama/Llama-2-7b-hf",
        "lora_rank": 8,
        "train_steps": 1000,
    }
)
artifact.add_dir("./model_output")   # 添加整个目录
artifact.add_file("./model_output/config.json")  # 添加单个文件
wandb.log_artifact(artifact)

# 方式2：记录单个文件
artifact = wandb.Artifact("tokenizer", type="model")
artifact.add_file("./tokenizer.model")
wandb.log_artifact(artifact)

# ===== 使用 Artifact =====
# 下载最新版本
artifact = wandb.use_artifact("llama2-lora-model:latest")
artifact_dir = artifact.download()  # 下载到本地并返回路径

# 下载特定版本
artifact = wandb.use_artifact("llama2-lora-model:v3")

# 下载特定别名的版本
artifact = wandb.use_artifact("llama2-lora-model:best-model")

# ===== 链接 Artifact 到运行 =====
# 在训练运行中记录Artifact时，W&B自动关联
# 也可以手动将输入/输出Artifact关联到运行
run = wandb.init()
input_artifact = run.use_artifact("dataset:v2")  # 标记为输入
output_artifact = wandb.Artifact("model", type="model")
output_artifact.add_dir("./checkpoints")
run.log_artifact(output_artifact)  # 标记为输出

# ===== 数据集版本管理 =====
dataset_artifact = wandb.Artifact(
    "alpaca-cleaned",
    type="dataset",
    metadata={"source": "tatsu-lab/alpaca", "split": "train", "size": 52002}
)
dataset_artifact.add_file("alpaca_train.jsonl")
wandb.log_artifact(dataset_artifact)

# ===== 步骤1：定义 Sweep 配置 =====
sweep_config = {
    "method": "bayes",  # 搜索策略: "grid", "random", "bayes"
    "metric": {
        "name": "eval/loss",
        "goal": "minimize"
    },
    "parameters": {
        "learning_rate": {
            "distribution": "log_uniform_values",
            "min": 1e-6,
            "max": 1e-3
        },
        "lora_rank": {
            "values": [4, 8, 16, 32]
        },
        "lora_alpha": {
            "distribution": "int_uniform",
            "min": 8,
            "max": 64
        },
        "batch_size": {
            "values": [8, 16, 32]
        },
        "weight_decay": {
            "distribution": "uniform",
            "min": 0.0,
            "max": 0.1
        },
        "warmup_steps": {
            "distribution": "int_uniform",
            "min": 50,
            "max": 500
        }
    },
    "early_terminate": {
        "type": "hyperband",
        "min_iter": 3,
        "s": 2
    },
    "run_cap": 20,  # 最大运行次数
}

# ===== 步骤2：创建 Sweep =====
sweep_id = wandb.sweep(
    sweep=sweep_config,
    project="llm-finetuning"
)
print(f"Sweep ID: {sweep_id}")

# ===== 步骤3：定义训练函数 =====
def train():
    # wandb.init() 会自动接收 Sweep 分配的超参数
    wandb.init()

    # 从 wandb.config 获取当前超参数组合
    config = wandb.config
    learning_rate = config.learning_rate
    lora_rank = config.lora_rank
    batch_size = config.batch_size

    # ... 训练逻辑 ...
    for step in range(100):
        loss = simulate_training(learning_rate, lora_rank, step)
        wandb.log({"train/loss": loss, "eval/loss": loss * 1.1})

    # 记录模型
    artifact = wandb.Artifact("model", type="model")
    artifact.add_file("./model.bin")
    wandb.log_artifact(artifact)

# ===== 步骤4：启动 Sweep Agent =====
wandb.agent(sweep_id, function=train, count=5)  # 运行5次

# ===== 也可以用命令行启动 =====
# wandb agent <entity>/<project>/<sweep_id>

from transformers import TrainingArguments, Trainer

# 方式1：通过 report_to 参数启用
training_args = TrainingArguments(
    output_dir="./results",
    num_train_epochs=3,
    per_device_train_batch_size=16,
    learning_rate=2e-5,
    weight_decay=0.01,
    logging_steps=10,
    eval_strategy="steps",
    eval_steps=100,
    save_strategy="steps",
    save_steps=500,
    report_to="wandb",  # 关键：将日志输出到W&B
    run_name="llama2-sft-run01",  # W&B中的运行名称
    load_best_model_at_end=True,
    metric_for_best_model="eval_loss",
)

# 方式2：同时报告到多个平台
training_args = TrainingArguments(
    ...
    report_to=["wandb", "tensorboard"],  # 同时输出到W&B和TensorBoard
)

# 完整训练示例
import wandb
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer
from datasets import load_dataset

# 初始化W&B
wandb.init(
    project="llm-finetuning",
    name="llama2-sft-alpaca",
    config={
        "model_name": "meta-llama/Llama-2-7b-hf",
        "dataset": "tatsu-lab/alpaca",
        "learning_rate": 2e-5,
        "epochs": 3,
        "batch_size": 16,
    }
)

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")
dataset = load_dataset("tatsu-lab/alpaca")

training_args = TrainingArguments(
    output_dir="./results",
    report_to="wandb",
    run_name=wandb.run.name,
    num_train_epochs=wandb.config.epochs,
    per_device_train_batch_size=wandb.config.batch_size,
    learning_rate=wandb.config.learning_rate,
    logging_steps=10,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset["train"],
    tokenizer=tokenizer,
)

trainer.train()

# 记录最终模型为Artifact
artifact = wandb.Artifact("llama2-sft-model", type="model")
artifact.add_dir("./results")
wandb.log_artifact(artifact)

wandb.finish()

# ===== wandb.Table — 记录结构化数据 =====
# 记录模型生成结果
columns = ["input_prompt", "generated_text", "reference", "score"]
data = [
    ["什么是机器学习？", "机器学习是AI的子领域...", "机器学习是一种...", 0.85],
    ["解释深度学习", "深度学习使用神经网络...", "深度学习是机器学习...", 0.78],
]
table = wandb.Table(columns=columns, data=data)
wandb.log({"generation_examples": table})

# 动态添加行
table = wandb.Table(columns=["step", "input", "output", "loss"])
for step, sample in enumerate(eval_samples):
    output = model.generate(sample["input"])
    table.add_data(step, sample["input"], output, sample["loss"])
wandb.log({"eval_results": table})

# ===== wandb.Image — 记录图像 =====
import matplotlib.pyplot as plt
import numpy as np

# 记录matplotlib图表
fig, ax = plt.subplots()
attention = np.random.rand(12, 12)
ax.imshow(attention, cmap="viridis")
ax.set_title("Attention Heatmap - Layer 6")
wandb.log({"attention_map": wandb.Image(fig)})
plt.close()

# 记录numpy数组为图像
image_array = np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8)
wandb.log({"sample_image": wandb.Image(image_array, caption="Generated sample")})

# ===== wandb.plot — 内置图表类型 =====
# 折线图
data = [[x, x**2] for x in range(100)]
table = wandb.Table(data=data, columns=["x", "y"])
wandb.log({"parabola": wandb.plot.line(table, "x", "y", title="Quadratic")})

# 柱状图
table = wandb.Table(
    data=[["train", 0.95], ["eval", 0.88], ["test", 0.87]],
    columns=["split", "accuracy"]
)
wandb.log({"accuracy_chart": wandb.plot.bar(table, "split", "accuracy")})

# 散点图
table = wandb.Table(data=[[i, i*0.5+np.random.randn()] for i in range(50)],
                    columns=["perplexity", "loss"])
wandb.log({"perplexity_vs_loss": wandb.plot.scatter(table, "perplexity", "loss")})

# 在训练异常时发送通知
import wandb

wandb.init(project="llm-finetuning")

for step in range(10000):
    loss = train_step()

    if loss > 10.0:
        wandb.alert(
            title="训练Loss爆炸！",
            text=f"Step {step}: loss = {loss:.4f}，超过阈值10.0，请检查学习率设置。",
            level=wandb.AlertLevel.WARN,  # INFO, WARN, ERROR
            wait_duration=300,  # 5分钟内不重复发送相同alert
        )

    if loss < 0.01:
        wandb.alert(
            title="训练收敛完成",
            text=f"Step {step}: loss = {loss:.6f}，训练已收敛。",
            level=wandb.AlertLevel.INFO,
        )

# Alert 会在 W&B 网页端显示，并可通过 Slack/Webhook 集成发送到其他渠道

# 方式1：通过环境变量设置
import os
os.environ["WANDB_MODE"] = "offline"  # 离线模式，日志保存在本地

# 方式2：通过 wandb.init 参数
wandb.init(mode="offline")

# 方式3：通过 wandb settings
wandb.init(settings=wandb.Settings(mode="offline"))

# 离线模式下，所有日志保存在本地 wandb/ 目录
# 训练完成后，使用命令行同步到云端：
# wandb sync ./wandb/offline-run-20231001_123456-abc123

# 批量同步所有离线运行
# wandb sync --sync-all

# 禁用W&B（完全不记录）
os.environ["WANDB_MODE"] = "disabled"
# 或
wandb.init(mode="disabled")

# 标记运行结束，上传所有未同步的数据
wandb.finish(exit_code=0)  # 0=成功，非0=失败

# 不调用finish()时，运行会在脚本退出时自动结束
# 但建议显式调用，特别是在Jupyter Notebook中

import wandb
from transformers import (
    AutoModelForCausalLM, AutoTokenizer,
    TrainingArguments, Trainer,
    DataCollatorForLanguageModeling
)
from peft import LoraConfig, get_peft_model
from datasets import load_dataset

def main():
    # 1. 初始化实验
    wandb.init(
        project="llm-instruction-tuning",
        name="mistral-7b-lora-alpaca",
        tags=["mistral", "lora", "sft"],
        config={
            "model_name": "mistralai/Mistral-7B-v0.1",
            "lora_rank": 16,
            "lora_alpha": 32,
            "lora_dropout": 0.05,
            "learning_rate": 3e-4,
            "epochs": 3,
            "batch_size": 4,
            "gradient_accumulation_steps": 8,
            "max_seq_length": 512,
            "dataset": "tatsu-lab/alpaca",
        }
    )
    config = wandb.config

    # 2. 加载模型与LoRA配置
    model = AutoModelForCausalLM.from_pretrained(
        config.model_name, torch_dtype=torch.bfloat16, device_map="auto"
    )
    lora_config = LoraConfig(
        r=config.lora_rank,
        lora_alpha=config.lora_alpha,
        lora_dropout=config.lora_dropout,
        target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
        task_type="CAUSAL_LM",
    )
    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()

    # 3. 记录数据集Artifact
    dataset_artifact = wandb.Artifact("alpaca-dataset", type="dataset")
    dataset_artifact.add_reference("https://huggingface.co/datasets/tatsu-lab/alpaca")
    wandb.log_artifact(dataset_artifact)

    # 4. 训练
    training_args = TrainingArguments(
        output_dir="./results",
        report_to="wandb",
        run_name=wandb.run.name,
        num_train_epochs=config.epochs,
        per_device_train_batch_size=config.batch_size,
        learning_rate=config.learning_rate,
        gradient_accumulation_steps=config.gradient_accumulation_steps,
        logging_steps=10,
        eval_strategy="steps",
        eval_steps=100,
        save_strategy="steps",
        save_steps=500,
        load_best_model_at_end=True,
        bf16=True,
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=tokenized_dataset["train"],
        eval_dataset=tokenized_dataset["test"],
    )
    trainer.train()

    # 5. 记录模型Artifact
    model_artifact = wandb.Artifact("mistral-lora-model", type="model")
    model_artifact.add_dir("./results/checkpoint-best")
    wandb.log_artifact(model_artifact)

    # 6. 记录生成样本
    table = wandb.Table(columns=["prompt", "generated", "reference"])
    for sample in eval_samples[:20]:
        generated = generate_text(model, sample["prompt"])
        table.add_data(sample["prompt"], generated, sample["output"])
    wandb.log({"generation_samples": table})

    wandb.finish()

if __name__ == "__main__":
    main()

import wandb

sweep_config = {
    "method": "bayes",
    "metric": {"name": "eval/loss", "goal": "minimize"},
    "early_terminate": {"type": "hyperband", "min_iter": 5, "s": 2},
    "parameters": {
        "learning_rate": {
            "distribution": "log_uniform_values",
            "min": 1e-5, "max": 5e-4
        },
        "lora_rank": {"values": [4, 8, 16, 32]},
        "lora_alpha": {
            "distribution": "int_uniform",
            "min": 8, "max": 64
        },
        "weight_decay": {
            "distribution": "uniform",
            "min": 0.0, "max": 0.1
        },
    },
}

sweep_id = wandb.sweep(sweep_config, project="llm-hp-search")

def sweep_train():
    wandb.init()
    config = wandb.config

    # 使用Sweep分配的超参数进行训练
    model = setup_model(lora_rank=config.lora_rank, lora_alpha=config.lora_alpha)
    train_result = train_model(
        model,
        learning_rate=config.learning_rate,
        weight_decay=config.weight_decay,
    )

    wandb.log({
        "eval/loss": train_result.eval_loss,
        "eval/perplexity": train_result.perplexity,
    })

wandb.agent(sweep_id, function=sweep_train, count=10)

import wandb

wandb.init(
    project="rlhf-training",
    name="llama2-7b-rlhf",
    config={
        "ppo_epochs": 4,
        "kl_coeff": 0.2,
        "reward_model": "reward-model-v2",
        "actor_lr": 1e-6,
        "critic_lr": 5e-6,
    }
)

# PPO训练循环
for ppo_epoch in range(num_ppo_epochs):
    for batch in ppo_dataloader:
        # 生成回复
        query_tensors, response_tensors = generate_responses(actor_model, batch)

        # 计算奖励
        rewards = compute_rewards(reward_model, query_tensors, response_tensors)

        # PPO更新
        stats = ppo_step(
            actor_model, critic_model,
            query_tensors, response_tensors, rewards
        )

        # 记录RLHF特有指标
        wandb.log({
            "rlhf/mean_reward": stats["mean_reward"],
            "rlhf/kl_divergence": stats["kl_divergence"],
            "rlhf/actor_loss": stats["actor_loss"],
            "rlhf/critic_loss": stats["critic_loss"],
            "rlhf/entropy": stats["entropy"],
            "rlhf/approx_kl": stats["approx_kl"],
            "rlhf/clip_fraction": stats["clip_fraction"],
            "rlhf/explained_variance": stats["explained_variance"],
        }, step=global_step)

        # 定期记录生成样本
        if global_step % 100 == 0:
            table = wandb.Table(columns=["prompt", "response", "reward"])
            for i in range(min(5, len(query_tensors))):
                prompt = tokenizer.decode(query_tensors[i])
                response = tokenizer.decode(response_tensors[i])
                table.add_data(prompt, response, rewards[i].item())
            wandb.log({"rlhf/samples": table}, step=global_step)

# 1. 控制log频率，避免过多小批次上传
# 不好：每步都log
for step in range(100000):
    wandb.log({"loss": loss}, step=step)  # 太频繁

# 好：每N步log一次
LOG_INTERVAL = 50
for step in range(100000):
    if step % LOG_INTERVAL == 0:
        wandb.log({"loss": loss}, step=step)

# 2. 使用 commit=False 合并同一步的多次log
wandb.log({"train/loss": loss}, commit=False)
wandb.log({"train/lr": lr}, commit=False)
wandb.log({"train/grad_norm": grad_norm}, commit=True)  # 最后一次commit

# 3. 大文件使用Artifact，不要用log
# 不好
wandb.log({"model_file": wandb.File("model.pt")]})  # 不适合大文件

# 好
artifact = wandb.Artifact("model", type="model")
artifact.add_file("model.pt")
wandb.log_artifact(artifact)

# 1. 使用命名约定
wandb.init(
    name=f"{model_name}-{method}-{dataset}-{timestamp}",
    # 如: "llama2-7b-lora-alpaca-20240101"
)

# 2. 善用tags进行筛选
wandb.init(tags=["lora", "v2", "production-candidate"])

# 3. 使用group组织Sweep实验
wandb.init(group="lora-rank-ablation")

# 4. 配置层次化
wandb.init(config={
    "model": {"name": "llama2-7b", "dtype": "bfloat16"},
    "training": {"lr": 2e-5, "epochs": 3, "batch_size": 16},
    "lora": {"rank": 8, "alpha": 16, "dropout": 0.05},
    "data": {"name": "alpaca", "split": "train", "size": 52002},
})

# 1. 使用entity指定团队
wandb.init(entity="my-research-team", project="llm-projects")

# 2. 使用报告(Reports)功能组织实验对比
# 在W&B UI中创建Report，嵌入多个运行的图表

# 3. 使用Alert通知关键节点
wandb.alert(title="训练完成", text="模型已收敛", level=wandb.AlertLevel.INFO)

# 4. 使用Artifact共享模型和数据
# 团队成员可通过 use_artifact 下载他人上传的Artifact
artifact = wandb.use_artifact("my-research-team/llm-projects/best-model:v5")

# 1. 敏感数据不要log
# 不好：记录用户原始输入
wandb.log({"user_input": sensitive_text})

# 好：记录脱敏后的统计信息
wandb.log({"input_length": len(text), "input_type": "question"})

# 2. 使用私有项目
# W&B默认项目为私有，团队内可见
# 需要公开时显式设置

# 3. API Key安全
# 不要将API Key硬编码在代码中
# 使用环境变量或 wandb login 命令