Krea 2 LoRA 训练全流程踩坑记录：从打标到双卡并行的 Windows 原生实战

双 RTX 16GB 显卡 + 机械硬盘 + Windows 10 原生环境（非 WSL2），完成 Krea 2 LoRA 训练的全流程实战。涵盖自动打标、显存优化、cuDNN NaN 逆向调试、双卡数据并行四大主题。所有踩坑均有源码级修复，最终实现单卡 3.6s/step、双卡 DDP 5.3s/step 的稳定训练。

一、环境概况

项目	配置
GPU	RTX 5060 Ti 16GB（sm_120）+ RTX 4060 Ti 16GB（sm_89）
系统	Windows 10，非 WSL2
训练框架	DiffPipeForge（diffusion-pipe + DeepSpeed + ComfyUI 子模块）
Python 环境	conda `dp271`：Python 3.11 + PyTorch 2.7.1+cu128 + DeepSpeed 0.18.4

为什么不用 WSL2

WSL2 的 drvfs 对机械盘读取速度仅 134 MB/s（NTFS 原生约 180 MB/s）。13GB 模型加载从 15 分钟变成 20+ 分钟，不值得。

为什么不用最新 PyTorch

PyTorch 2.11（最新版）在这个环境下每步训练 18 分钟，且 DeepSpeed checkpoint 保存时 segfault。便携版 DiffPipeForge 用的是 PyTorch 2.7.1，经过社区验证稳定。最终完全对齐便携版版本搭建独立 conda 环境，摆脱便携版依赖。

机械盘的隐性代价

13GB 模型从 HDD 加载约 5-10 分钟，Joy Caption 16GB 模型约 2.5 分钟。模型加载是一次性的，推理/训练阶段全在 GPU 上，不受磁盘影响。但 ComfyUI 每次生成后会卸载模型，重新加载同样慢------测试 LoRA 效果时需注意。

二、自动打标：Joy Caption Alpha Two

为什么需要描述性 caption

初始训练只用触发词作为 caption，1 epoch 生成的 LoRA 效果与预期差异大。分析后发现两个问题：

训练量不足：单 epoch 曝光次数仅为最低推荐量（1500-3000）的 35%
caption 过简单：全量触发词让模型无法学习图像的视觉特征描述，LoRA 学到的更多是"整体风格漂移"而非"精准的人物/场景还原"

改用 Joy Caption Alpha Two（Llama 3.1 8B Instruct + SigLIP 视觉编码器，8.48B 参数）自动生成描述性 caption。

模型信息

项目	值
架构	LlavaForConditionalGeneration
参数量	8.48B，bf16，15.81GB（4 分片）
视觉编码器	SigLIP2 so400m patch14-384
要求	transformers >= 4.51（实际用 4.56.2）

踩坑 1：17GB 模型 vs 16GB 显存

Joy Caption bf16 实际占 ~16.1GB，单卡 16GB 放不下。

现象：命令行设了 CUDA_VISIBLE_DEVICES=0（只暴露一张卡）时勉强塞进（16097/16384 MiB，剩 287MiB）；但交互模式下双卡都可见，PyTorch 初始化双卡 context 多占显存，第 4 个分片加载时差 112MiB → OOM。

解决：在加载模型前设 CUDA_VISIBLE_DEVICES 只暴露一张卡，省掉多余的 context 开销：

python 复制代码

os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_id)
# 此后 CUDA 只看到一张卡，编号为 0
llava_model = LlavaForConditionalGeneration.from_pretrained(
    MODEL_PATH, torch_dtype=torch.bfloat16, device_map=0,
)

也尝试过 bitsandbytes NF4 int4 量化（17GB → ~5GB），但在线量化加载非常慢，不推荐。bf16 + 单卡隔离是更好的选择。

踩坑 2：生成陷入死循环

现象：do_sample=True 时偶尔生成重复 token，一张图推理超过 10 分钟不结束。第一张图约 2 分钟完成且质量很好，第二张就卡死了，没有任何报错。

解决：加 repetition_penalty=1.1，之后再未出现循环：

python 复制代码

generate_ids = llava_model.generate(
    **inputs,
    max_new_tokens=150,
    do_sample=True,
    temperature=0.6,
    top_p=0.9,
    repetition_penalty=1.1,  # 关键修复
    use_cache=True,
)

踩坑 3：Llava chat template 脆弱性

HuggingFace 官方 README 特别警告：Llava 模型的 chat 处理很脆弱，错误的组合会导致多个 <bos> token，模型效果骤降。必须严格按官方示例的组合调用：

python 复制代码

# 正确组合
convo_string = processor.apply_chat_template(convo, tokenize=False, add_generation_prompt=True)
inputs = processor(text=[convo_string], images=[image], return_tensors="pt").to('cuda')
inputs['pixel_values'] = inputs['pixel_values'].to(torch.bfloat16)

打标效果

erlang 复制代码

速度：~100 秒/张
显存：16097 MiB / 16384 MiB（99%）

caption 质量示例（示例图片为虚构，展示输出格式）：

photograph of a woman with long brown hair, standing outdoors in a garden, wearing a white summer dress with floral patterns, holding a book in her hands, smiling expression, soft afternoon sunlight, shallow depth of field, realistic photography style, natural colors, <trigger_word>

准确覆盖人物外貌、服装、姿势、表情、光线、场景、风格，末尾自动追加触发词。

三、16GB 显存如何训练 24GB 模型

问题

Krea 2 原始 bf16 权重 24.76GB，单卡 16GB 根本放不下。

尝试过的方案

方案	结果	原因
block swap（GPU↔CPU 权重交换）	❌ OOM	`enable_block_swap` 初始化时先把整个模型 `.to('cuda')`，16GB 直接爆
双卡流水线并行	❌ 死锁	后文第五节详述
运行时降 fp8	❌ 性能差	每次 forward 都要 dequantize

最终方案：fp8 预量化 + activation checkpointing

1. 使用预量化的 fp8 模型（13GB）

直接使用 krea2_turbo_fp8.safetensors（12.9GB），配置中指定 diffusion_model_dtype = 'float8'：

toml 复制代码

[model]
diffusion_model = 'krea2_turbo_fp8.safetensors'
diffusion_model_dtype = 'float8'  # 关键：不设则默认 dequantize 到 bf16，显存翻倍 OOM

2. 开启 activation checkpointing

toml 复制代码

activation_checkpointing = true  # 省约 6GB 激活内存，代价是 ~20% 速度

关闭时峰值显存 21.96GB > 16GB，开启后降到 ~13GB。

组合效果：16GB 显存从"放不下"变成"有 3GB 余量"。

教训

16GB 显存做大模型 LoRA 训练没有银弹。fp8 预量化省模型权重内存，activation checkpointing 省激活内存，两者必须同时上。block swap 的"先把整模型搬上 GPU 再逐步 swap 出来"初始化机制，在低显存场景下反而是陷阱。

四、🔥 NaN 梯度（核心问题）

现象

Step 1 loss 完全正常（0.3~0.5），但 optimizer.step() 执行后 LoRA 参数变成 NaN。从 step 2 起 loss 永久为 NaN。

ini 复制代码

step 1: loss=0.35  ✅
step 2: loss=nan   ❌（之后再也没恢复过）

排查过程

第一步：排除常见原因

排查项	结论
优化器类型（AdamW / adamw_optimi）	❌ 无关，两种都 NaN
GPU 架构（sm_120 / sm_89）	❌ 无关，两张卡都 NaN
gradient_clipping	❌ 无关
text encoder 精度（fp8 / bf16）	❌ 无关
数据集缓存值（latents 统计）	❌ 无异常
learning rate 过高	❌ 无关，降到极小仍 NaN

所有常规原因全部排除。基本确定：问题出在计算图本身，不是超参数。

第二步：梯度 Hook --- 缩小范围到参数级别

在所有 LoRA 参数上注册 register_hook，在 zero_grad() 清零前捕获真实梯度值：

python 复制代码

def hook(grad):
    g_nan = torch.isnan(grad).any().item()
    g_max = grad.abs().max().item() if not g_nan else float('inf')
    _grad_debug['data'][name] = (g_nan, g_inf, g_max)
    return grad

for name, p in model.named_parameters():
    if p.requires_grad:
        p.register_hook(hook)

发现了关键线索------只有 query projection 的梯度是 NaN：

ini 复制代码

wo.lora_B:   ✅ g_absmax=1.13e-03
wk.lora_B:   ✅ g_absmax=3.23e-03
wv.lora_B:   ✅ g_absmax=1.66e-03
gate.lora_B: ✅ g_absmax=4.48e-04
wq.lora_B:   ❌ g_nan=True    ← 只有这里！
wq.lora_A:   ❌ g_nan=True    ← 只有这里！

其他所有投影（output、key、value、gate）的梯度完全正常。如果是数据问题或超参数问题，不应该只影响一个投影。

第三步：Backward Hook --- 缩小范围到层级

Krea 2 使用 per-head QK-Norm（对 Q 和 K 分别做 RMSNorm）。在 qnorm 和 knorm 的输出上注册 backward hook，追踪梯度从 attention 传回时的状态：

python 复制代码

class QKNorm(nn.Module):
    def forward(self, q, k):
        qn = self.qnorm(q)
        kn = self.knorm(k)

        def _mk_hook(name):
            def hook(grad):
                isnan = torch.isnan(grad).any().item()
                gmax = grad.abs().max().item() if not isnan else float('inf')
                print(f'[BKWD layer{idx} {name}] nan={isnan} absmax={gmax}')
                return grad
            return hook

        qn.register_hook(_mk_hook('q_grad_from_attn'))
        kn.register_hook(_mk_hook('k_grad_from_attn'))
        return qn, kn

结果揭示了 NaN 的精确起源------layer 5 的 Q 梯度：

r 复制代码

layer 32→6:  q_grad ✅  k_grad ✅     (全部正常)
layer 5:     q_grad ❌  k_grad ✅     (4.6e-03，正常！NaN 唯一起点)
layer 4→1:   q_grad ❌  k_grad ❌     (NaN 逐层传播)

关键发现：在 layer 5，Q 和 K 走的是完全相同的代码路径（相同的 RMSNorm 结构、相同的 attention 函数），但只有 Q 的梯度变 NaN，K 完全正常。

这说明问题不在 RMSNorm，而在 attention 计算的反向传播中，对 Q 的梯度计算出了问题 。layer 5 可能只是恰好触发了某种极端数值条件（如 Q·K^T 的某些值过大，导致 softmax backward 产生 0 × ∞ = NaN）。

根因

追踪到 comfy/ops.py 中的 attention 后端选择逻辑：

python 复制代码

# comfy/ops.py --- Windows 上强制优先使用 cuDNN attention
if torch.cuda.is_available() and WINDOWS:
    SDPA_BACKEND_PRIORITY = [
        SDPBackend.CUDNN_ATTENTION,     # ← 最高优先级
        SDPBackend.FLASH_ATTENTION,
        SDPBackend.EFFICIENT_ATTENTION,
        SDPBackend.MATH,
    ]

ComfyUI 在 Windows 上强制优先使用 cuDNN attention 后端。 该后端在 bf16 反向传播中，对 Q 的梯度计算存在数值 bug------在某些层的特定数值条件下产生 NaN，然后通过反向传播逐层扩散，最终污染所有参数。

踩过的弯路

尝试	为什么失败
Q/K/V 提升到 float32 再做 attention	❌ 无效。`comfy.ops` 内部的 `sdpa_kernel` 仍选择 cuDNN 后端
在 model.py 外层包 `sdpa_kernel(MATH)`	❌ 无效。`comfy.ops` 内部嵌套覆盖了外层设置

最终修复

直接修改 comfy/ops.py，强制只使用 MATH 后端：

python 复制代码

# comfy/ops.py --- 修复后
def scaled_dot_product_attention(q, k, v, *args, **kwargs):
    with sdpa_kernel([SDPBackend.MATH], set_priority=True):
        return torch.nn.functional.scaled_dot_product_attention(q, k, v, *args, **kwargs)

修复后训练稳定，loss 从 0.35 下降到 0.1-0.19 范围。代价：每步从 3.3s 涨到 3.6s（MATH 比 cuDNN 慢约 10%）。

调试方法论

r 复制代码

参数级 hook → 发现只有 wq 梯度 NaN（其他投影正常）
    ↓
层级 backward hook → 发现只有 layer 5 的 Q 梯度 NaN（K 正常）
    ↓
源码追踪 → 发现 cuDNN attention 后端的 backward bug
    ↓
强制 MATH 后端 → 问题解决

相比盲目改超参数，这种"梯度 hook 定位法"能精确定位问题源头。

五、双卡数据并行

为什么流水线并行失败

最先尝试 pipeline parallel（每卡跑一半 transformer 层）：

toml 复制代码

pipeline_stages = 2
partition_method = 'manual'
partition_split = [15]  # layer 0-14 在 GPU 0，layer 15-29 在 GPU 1

结果：deepspeed.initialize() 永远不返回，死锁。

根因：pipeline parallel 需要 GPU 之间做 point-to-point 通信（前向传 activation，反向传 gradient）。Windows 上 DeepSpeed 只能用 Gloo 后端（没有 NCCL），而 Gloo 的 point-to-point 通信在 Windows 上有已知问题。

数据并行的三个障碍

转向 data parallel（每卡跑完整模型，处理不同 batch，梯度 allreduce 同步）。这条路只需要 allreduce（集合通信），Gloo 支持。但 diffusion-pipe 的代码有三个为单卡硬编码的地方：

障碍 1：distributed_init 硬编码 world_size=1

python 复制代码

# 修复前
def distributed_init(args):
    world_size = 1          # ← 硬编码
    return world_size, rank, local_rank

# 修复后
def distributed_init(args):
    world_size = int(os.environ.get('WORLD_SIZE', '1'))   # 从环境变量读取
    rank = int(os.environ.get('RANK', '0'))
    local_rank = int(os.environ.get('LOCAL_RANK', '0'))
    return world_size, rank, local_rank

障碍 2：init_method 使用 PID-based 文件路径

python 复制代码

# 修复前 --- 每个 rank 用不同的文件，无法发现彼此
init_file = os.path.join(tempfile.gettempdir(), f'deepspeed_init_{os.getpid()}')
init_method = f'file://{init_file}'

# 修复后 --- 用环境变量方式，所有 rank 共享 MASTER_ADDR:PORT
deepspeed.init_distributed(dist_backend='gloo', init_method='env://',
                           rank=rank, world_size=world_size)

障碍 3：主机名被解析为 Docker 内部地址

css 复制代码

[hostname 07:26:55] The client socket has failed to connect to
[kubernetes.docker.internal]:29500 (system error: 10049)

Windows 上的 Docker Desktop 会修改 DNS 解析，把本机主机名映射到 kubernetes.docker.internal，导致进程间无法通信。

python 复制代码

# 修复 --- 强制使用 loopback 地址
os.environ['MASTER_ADDR'] = '127.0.0.1'

数据并行结果

指标	单卡	双卡 DDP	提升
Global batch size	1	2	2×
每步耗时	3.6s	5.3s	+47%
samples/sec	0.27	0.38	+41%
每 epoch（~500 图）	~32 分钟	~23 分钟	省 9 分钟

没有实现理论上的 2× 提速，因为：

两张卡模型不同（sm_120 vs sm_89），速度有差异
Gloo allreduce 的梯度同步开销（每步额外 ~1.7s）

六、最终方案

训练配置

toml 复制代码

[model]
type = 'krea2'
diffusion_model = 'krea2_turbo_fp8.safetensors'   # fp8 预量化（13GB）
diffusion_model_dtype = 'float8'
text_encoders = [{ path = 'qwen3vl_4b_bf16.safetensors', type = 'krea2' }]

activation_checkpointing = true    # 必须开
micro_batch_size_per_gpu = 1
gradient_accumulation_steps = 1
gradient_clipping = 1.0

[optimizer]
type = 'adamw_optimi'
lr = 2e-5

[adapter]
type = 'lora'
rank = 32
dtype = 'bfloat16'

代码修改汇总（共 4 处）

文件	修改	解决问题
`comfy/ops.py`	强制 `SDPBackend.MATH`	NaN 梯度（cuDNN bf16 backward bug）
`train.py` `distributed_init`	从环境变量读取 world_size/rank	单卡硬编码
`train.py` `init_distributed`	`env://` + `MASTER_ADDR=127.0.0.1`	Windows Docker DNS
配置层面	fp8 模型 + activation checkpointing	16GB 显存不足

启动命令

bash 复制代码

# 环境变量
set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
set DS_BUILD_AIO=0
set DS_BUILD_CUFILE=0
set MASTER_ADDR=127.0.0.1
set NCCL_P2P_DISABLE=1
set NCCL_IB_DISABLE=1

# 单卡
python -m deepspeed.launcher.runner --num_gpus=1 train.py --deepspeed --config krea2_train.toml

# 双卡 DDP
python -m deepspeed.launcher.runner --num_gpus=2 train.py --deepspeed --config krea2_train.toml

端到端工作流

markdown 复制代码

1. Joy Caption 打标（~100s/张）
   ↓ caption 变了，需重新缓存
2. 预缓存数据集（VAE latents + text embeddings）
   ↓
3. 训练（单卡 3.6s/step 或双卡 5.3s/step）
   ↓ 每 epoch 自动保存 LoRA
4. 复制 LoRA 到 ComfyUI，测试生成效果
   ↓ strength 0.3-0.5 起步

七、经验总结

caption 质量决定 LoRA 效果上限。 纯触发词打标让模型只学到"整体漂移"。Joy Caption 生成的描述性 caption（外貌/服装/姿势/场景）让模型学到精准的视觉特征映射。17GB 模型塞 16GB 卡用 CUDA_VISIBLE_DEVICES 单卡隔离，比量化更简单。
NaN 调试：梯度 Hook 定位法。 与其盲目调超参数，不如在参数和层级别上注册 backward hook，精确追踪 NaN 的起源。这次调试从"所有投影"缩小到"只有 wq"，再缩小到"只有 layer 5 的 Q 梯度"，最终定位到 cuDNN 后端。
cuDNN attention 的 bf16 backward bug。 任何在 Windows + bf16 + cuDNN attention 后端上训练的扩散模型都可能遇到。如果你看到"只有 query projection 梯度 NaN"的症状，第一反应应该是检查 attention 后端，强制切到 MATH。
Windows 上 DeepSpeed 的分布式限制。 没有 NCCL，只能用 Gloo。Gloo 支持 allreduce（数据并行可用），不支持 pipeline point-to-point（流水线并行死锁）。Docker Desktop 会污染 DNS 解析，必须强制 MASTER_ADDR=127.0.0.1。
16GB 显存的生存策略。 fp8 预量化 + activation checkpointing 是必需组合。block swap 的初始化机制不适合低显存场景。
do_sample 生成加 repetition_penalty。 VLM 在采样模式下偶尔陷入重复 token 死循环，无报错无超时。加 1.1 的惩罚项即可避免。

环境：Windows 10 / RTX 5060 Ti 16GB + RTX 4060 Ti 16GB / PyTorch 2.7.1+cu128 / DeepSpeed 0.18.4 / transformers 4.56.2