litgpt Thunder 如何注册 Unsloth Triton 核函数作为自定义 executor【免费下载链接】litgpt20 high-performance LLMs with recipes to pretrain, finetune and deploy at scale.项目地址: https://gitcode.com/GitHub_Trending/li/litgpt用 Lightning Thunder 给 litgpt 的GPT模型做 JIT 编译时默认所有算子都由 Thunder 内置 executor如 NvFuser处理。如果你希望把 cross entropy loss、SwiGLU 激活、RoPE 位置编码这几个算子替换成 Unsloth 团队手写的 Triton 核函数做法是把它们注册成一个 Thunder 自定义 executor再把该 executor 放进thunder.jit的executors优先级列表。litgpt 仓库在 extensions/thunder/unsloth/ 目录已经内置了这个 executor本文基于仓库文档与源码说明它的注册方式、启用路径和验证方法。前提条件均来自仓库文档CUDA GPU。executor 里每个算子的 checker 都要求输入张量在cuda设备上非 CUDA 设备不会触发替换环境能导入thunder和triton。litgpt/constants.py 中用RequirementCache(thunder)和RequirementCache(triton)探测这两个可选依赖kernels 代码只在 triton 可用时才import tritonpyproject.toml 要求 Python 3.10并把lightning-thunder0.2.dev20250119限定sys_platformlinux且 python 3.10放在可选依赖组compiler中。按 pyproject 的可选依赖组安装即pip install litgpt[compiler]在 litgpt 仓库根目录下运行后续导入路径都以根目录为基准。需要提醒extensions/thunder/README.md 明确标注这是 early-access 开发版、仅供内部使用并说明 Lightning Thunder 处于 alpha 阶段、不适用于生产任务最新用法建议直接查阅 Lightning Thunder 项目本身。最短路径启用仓库自带的 unsloth executorThunder 的thunder.jit接受一个 executor 优先级列表靠前的 executor 优先认领算子。extensions/thunder/README.md 指出顺序很重要unsloth 必须排在nvfuser之前这样自定义算子会在 NvFuser 创建融合区域之前执行import thunder model thunder.jit( model, executors[sdpa, unsloth, torchcompile_cat, nvfuser, torch] )列表里的字符串unsloth只有在 executor 模块被导入、注册动作执行完之后才有效。导入即注册from extensions.thunder.unsloth.executor import unsloth_ex # import for registrationextensions/thunder/unsloth/executor.py 共注册三类算子算子对应 litgpt 中的实现注册的 unsloth 算子Cross entropy lossltorch.cross_entropyunsloth_cross_entropy/unsloth_cross_entropy_backwardSwiGLULLaMAMLP的激活litgpt.model.LLaMAMLPlitgpt_swiglu→unsloth_swiglu_forward/unsloth_swiglu_backwardRoPElitgpt.model.apply_ropeunsloth_apply_rope/unsloth_apply_rope_backwardTriton 核函数本体在 extensions/thunder/unsloth/kernels/ 下cross_entropy_loss.py、swiglu.py、rope_embedding.py。注册机制拆解从 executor.py 看如何注册一个核函数以下以 cross entropy 为例说明 executor.py 的完整注册套路写自己的 executor 时可以照搬这个结构。1. 创建 executor 并注册到 Thunderexecutor.py#L24-L25from thunder.extend import OperatorExecutor, register_executor unsloth_ex OperatorExecutor(unsloth, version0.1) register_executor(unsloth_ex)register_executor执行后executors列表里就可以使用字符串unsloth引用它也可以直接传 executor 对象仓库测试用的就是对象形式。2. 为每个核函数写 meta 函数并注册算子。meta 函数在 trace 阶段用TensorProxy描述输出张量的形状、dtype、设备和requires_gradfn是实际调用 Triton 核函数的实现executor.py#L48-L50unsloth_cross_entropy unsloth_ex.register_operator( unsloth_cross_entropy, metaunsloth_cross_entropy_meta, fnkernels.cross_entropy_loss._cross_entropy_forward_impl )注意unsloth_cross_entropy_meta里有一行注释the cross entropy kernel only supports float32即 loss 和 logsumexp 的输出 dtype 固定为float32。3. 注册对原算子的实现映射executor.py#L132-L137。这是让 Thunder 把 PyTorch 算子改写为 unsloth 算子的关键一步unsloth_ex.register_implementation( ltorch.cross_entropy, checkerunsloth_cross_entropy_checker, execution_transformlambda *args: cross_entropy_to_unsloth(*args)[0], grad_transformunsloth_cross_entropy_grad, )三个回调的职责checker决定本次调用是否走 unsloth 路径。cross entropy 的条件executor.py#L69-L88weight/size_average/reduce为Nonereduction是none或meanignore_index -100label_smoothing 0.0且 logits 与 labels 都在 cuda 上。任一条件不满足Thunder 会退回对该算子的常规分解而不是报错execution_transform生成前向实现。其中meanreduction 不在核函数内部代码用true_divide(sum(loss), n_items)在核外计算n_items sum(ne(labels, -100))源码 TODO 注释指出没有考虑所有元素都被 mask 时可能除零的情况grad_transform生成反向实现调用unsloth_cross_entropy_backward。反向实现里对 logits 做了clone()因为核函数会原地写梯度。SwiGLU 与 RoPE 的注册用了另一种变体replaces。这两个算子原本不是标准 PyTorch 算子而是 litgpt 自己模块里的调用所以 executor 先把入口函数替换成可被 Thunder 跟踪的自定义算子SwiGLUexecutor 在模块底部执行litgpt.model.LLaMAMLP ThunderLLaMAMLPexecutor.py#L157-L176子类化的forward把激活改写成swiglu(e, g)即torch.nn.functional.silu(e) * g再用register_operator(litgpt_swiglu, ..., replacesswiglu)让 Thunder 把它当作可替换的自定义算子随后register_implementation把它映射到kernels.swiglu_fg_kernel前向和kernels.swiglu_DWf_DW_dfg_kernel反向checker 只要求输入在 cuda 上RoPEregister_operator(litgpt_apply_rope, ..., replaceslitgpt.model.apply_rope)executor.py#L234-L236替换原函数checker 要求输入是 4D 且三个张量都在 cuda 上。反向核函数需要前向缓存的n_groups、BLOCK_SIZE、num_warps这些值由 meta 函数一并算出executor.py#L239-L247块大小由 kernels/utils.py 的calculate_settings决定上限MAX_FUSED_SIZE 65536。一个明确的边界RMSNorm没有集成 unsloth 核函数executor.py#L140-L147 的注释说明原因是核函数结果与 PyTorch 实现数值不等价且没有计算 weight 的梯度。接入 pretrain 脚本并运行extensions/thunder/pretrain.py 是 README 提供的、已接入该 executor 的训练脚本副本脚本的jit()pretrain.py#L508-L514先from unsloth.executor import unsloth_ex # import for registration这里靠脚本开头sys.path.append了extensions/thunder目录所以是unsloth.executor而不是extensions.thunder.unsloth.executor再执行thunder.jit(fn, executorsexecutors)。也就是说只要脚本跑起来unsloth executor 就已经注册你只需在命令行把unsloth加进 executor 列表前向损失函数用chunked_cross_entropy(logits, targets, chunk_size0)pretrain.py#L49-L53源码注释写明这是为了启用 unsloth cross entropy 核函数chunk 切分会改变调用形式导致 checker 不命中。README 给出的完整复现命令单卡# 下载 TinyLlama tokenizer仅 tokenizer litgpt download --repo_id TinyLlama/TinyLlama-1.1B-Chat-v1.0 --tokenizer_only true # 启用 unsloth executor 的 Thunder 预训练 python extensions/thunder/pretrain.py --config config.yaml \ --executors [sdpa, unsloth, torchcompile_cat, nvfuser, torch] --devices 1--config config.yaml的内容来自 README 的复现说明out_dir: out/pretrain-thunder data: TinyStories tokenizer_dir: checkpoints/TinyLlama/TinyLlama-1.1B-Chat-v1.0 logger_name: csv去掉unsloth的对照命令用于对比python extensions/thunder/pretrain.py --config config.yaml \ --executors [sdpa, torchcompile_cat, nvfuser, torch] --devices 1多卡时脚本默认走ThunderFSDPStrategystrategy参数默认fsdp。验证 executor 是否生效仓库测试 tests/ext_thunder/test_unsloth_executor.py 提供了两层验证trace 中出现 unsloth 算子名 数值与 eager 实现一致。下面这个最小片段取自其中的test_unsloth_cross_entropy要求 1 块 CUDA GPU 和 thunderimport torch import thunder from extensions.thunder.unsloth.executor import unsloth_ex logits torch.randn(64, 128, devicecuda, requires_gradTrue) labels torch.randint(128, (64,), devicecuda) def foo(logits, labels): return torch.nn.functional.cross_entropy(logits, labels, reductionmean, ignore_index-100) cfoo thunder.jit(foo, executors[unsloth_ex]) actual cfoo(logits, labels) # 前向 trace 应出现 unsloth 算子且不应混入 backward 算子 trace_str str(thunder.last_traces(cfoo)[-1]) assert unsloth_cross_entropy in trace_str and backward not in trace_str bwd_str str(thunder.last_backward_traces(cfoo)[-1]) assert unsloth_cross_entropy_backward in bwd_str # 数值与梯度都应和未编译的 eager 版本一致 expected foo(logits, labels) torch.testing.assert_close(actual, expected)整模型级别test_unsloth_gpt用一个两层小模型Config(vocab_size320, n_layer2, n_head4, n_embd64, ...)编译forward_and_loss内部同样是chunked_cross_entropy(..., chunk_size0)然后断言前向 trace 同时包含unsloth_cross_entropy、unsloth_apply_rope、unsloth_swiglu反向 trace 包含对应的三个*_backward算子。README 中也展示了一段开启 executor 后的前向/反向 trace 摘录文档示例可以看到算子已被改写为(t121, _, _, _, _, _) unsloth_apply_rope(t120, t21, t22) (t189, t190) unsloth_cross_entropy(t187, t188) # backward: t652 unsloth_cross_entropy_backward(t651, t187, t188, t190) t763 unsloth_apply_rope_backward(t757, t21, t22, 1, 8, 4)跑通 pretrain 脚本后观察终端按log_interval打印的loss train是否正常下降、无算子回退报错也是文档给出的运行路径的一部分。限制与已知问题性能结论README 的基准表文档数据NVIDIA A100-SXM4-40GB、TinyLlama 1.1B、TinyStories、单卡、step 10 的 ms/iter显示sdpa, torchcompile_cat, nvfuser, torch组合为 322.25 ms/iter、27.42 GB把unsloth加进列表后为 331.92 ms/iter、25.19 GB。README 据此认为这套手写核函数“似乎不值得”these hand-written kernels do not seem to be worth itNvFuser 的自动融合在该场景下更快。内存占用更低是唯一占优的指标meanreduction 的边界核外除法没有处理所有 label 都被-100mask 的情况可能除零源码 TODOdtype 限制cross entropy 核函数输出的 loss/logsumexp 只有 float32RMSNorm 未接入原因见上文注释预训练脚本中禁用了梯度裁剪fabric.clip_gradients一行被注释掉注释为 THUNDER unsupported链接到 Lightning Thunder 的 issue脚本不支持--compiler torch即不用 Thunder 的torch.compile原因是它无法编译_FabricModuleREADME 指向 PyTorch 的对应 issueexecutor 与 Thunder 版本绑定pyproject.toml 中compiler组要求lightning-thunder0.2.dev20250119而 README 基准数据来自一套 2024 年的开发版环境lightning-thunder0.2.0.dev20240505等两个文档中的版本不一致升级 Thunder 后应以 trace 验证结果为准。下一步如果要继续扩展比如替换其他算子或调整融合策略README 建议直接以 Lightning Thunder 项目文档为准本目录的内容作为 litgpt 侧的接入示例参考。【免费下载链接】litgpt20 high-performance LLMs with recipes to pretrain, finetune and deploy at scale.项目地址: https://gitcode.com/GitHub_Trending/li/litgpt创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
