tinygrad 环境变量完全指南：DEV、DEBUG、BEAM 等运行时开关的用法与源码级解读

tinygrad 环境变量完全指南DEV、DEBUG、BEAM 等运行时开关的用法与源码级解读【免费下载链接】tinygradYou like pytorch? You like micrograd? You love tinygrad! ❤️项目地址: https://gitcode.com/GitHub_Trending/tiny/tinygradtinygrad 将绝大多数运行时行为开关设计为环境变量从选择后端设备DEV、控制调试输出粒度DEBUG、调整内核搜索策略BEAM到切换默认浮点精度DEFAULT_FLOAT均可在不修改任何代码的情况下通过一行环境变量生效。本文以官方文档 docs/env_vars.md 为骨架结合 tinygrad 源码tinygrad/helpers.py、tinygrad/device.py、tinygrad/engine/realize.py 等逐项讲解每个变量的取值、作用与底层实现读完即可熟练使用DEVAMD:LLVM DEBUG4这类组合进行设备选择、性能分析和内核调试。环境变量的两种生效方式tinygrad 的环境变量既可以在进程启动前通过 shell 设置也可以在代码运行期内临时切换。官方文档给出的进程级用法是DEVCL DEBUG4 python3 -m pytest即在命令行前缀中一次性设置多个变量覆盖整个进程的运行时行为。用Context装饰器限定单个函数对于 tinygrad 开发者可以在函数上使用Context(...)装饰器让某个环境变量只在函数体内生效# in tensor.py (probably only useful if you are a tinygrad developer) Context(DEBUG4) def numpy(self) - ...Context在源码中继承自contextlib.ContextDecorator实现于 tinygrad/helpers.py。它进入作用域时把旧值暂存到old_context退出时逐一恢复因此可以安全嵌套使用def __enter__(self): self.old_context:dict[str, Any] {k: ContextVar._cache[k].value for k in self.kwargs} for k,v in self.kwargs.items(): ContextVar._cache[k].value v def __exit__(self, *args): for k,v in self.old_context.items(): ContextVar._cache[k].value v用with Context(...)限定代码块更常见的做法是配合with语句临时调整某个作用域内的行为with Context(DEBUG0): a Tensor.ones(10, 10) a * 2从源码看所有可配置变量都是ContextVar的实例定义于 tinygrad/helpers.py。ContextVar.__init__会调用getenv(key, default_value)读取进程环境变量作为初始值因此环境变量是全局默认值而Context是局部覆盖二者共享同一套变量名。getenv的实现也值得一提tinygrad/helpers.pyfunctools.cache def getenv(key:str, default:Any0): return type(default)(os.getenv(key, default))它会将字符串环境变量强制转换为默认值对应的类型如int、str、float并带functools.cache缓存这解释了为什么文档表格中的整数值可以直接参与比较运算。全局变量表控制核心运行时行为的开关以下变量控制 tinygrad 作为库被使用时import tinygrad之后的核心行为。表中#表示该变量可取任意整数值。表格内容完整继承自 docs/env_vars.md并补充了源码中的默认值与实现位置。VariablePossible Value(s)DescriptionDEBUG[1-7]enable debugging output (operations, timings, speed, generated code and more)DEV[AMD, NV, ...]enable a specific backend, see DEV 变量BEAM[#]number of beams in kernel beam searchDEFAULT_FLOAT[HALF, ...]specify the default float dtype (FLOAT32, HALF, BFLOAT16, FLOAT64, ...), default to FLOAT32IMAGE[1]enable 2d specific optimizationsFLOAT16[1]use float16 for images instead of float32JIT[0-2]0disabled, 1jit enabled (default), 2jit enabled, but graphs are disabledVIZ[1]0disabled, 1viz enabledALLOW_TF32[1]enable TensorFloat-32 tensor cores on Ampere or newer GPUs.WEBGPU_BACKEND[WGPUBackendType_Metal, ...]Force select a backend for WebGPU (Metal, DirectX, OpenGL, Vulkan...)CUDA_PATHstrUseCUDA_PATH/includefor CUDA headers for CUDA and NV backends. If not set, TinyGrad will use/usr/local/cuda/include,/usr/includeand/opt/cuda/include.各变量的源码默认值与实现位置从 tinygrad/helpers.py 可以看到这些变量的实际定义及其默认值DEV_DEV(DEV, )默认空字符串表示自动选择设备DEBUGContextVar(DEBUG, 0)默认关闭调试BEAMContextVar(BEAM, 0)默认不做 beam searchNOOPTContextVar(NOOPT, 0)。IMAGE、FLOAT16默认均为 0JIT默认为 1启用 JITJIT_BATCH_SIZE默认为 32。DEFAULT_FLOAT默认为float32DEFAULT_INT默认为int32二者由 dtype 系统在创建张量时读取。VIZ默认为 0源码注释指出 VIZ implies PROFILE, but you can run PROFILE without VIZ即PROFILE ContextVar(PROFILE, abs(VIZ.value))会随 VIZ 自动开启性能采样。ALLOW_TF32默认为 0tinygrad/helpers.py源码注释明确其用途是 allow tf32 to be used on NVIDIA GPUs。文档未列出但同处定义的CCACHE默认 1编译器缓存与SCACHE默认 1调度器缓存也可用相同方式在命令行覆盖。几个易混淆变量的实际作用JIT0/1/2默认 1 启用 JIT 编译将多次 realize 合并为一次内核启动设为 2 时 JIT 仍开启但禁用计算图graph记录通常用于排查 graph 相关行为设为 0 则完全关闭 JIT每次 realize 立即执行。IMAGE 与 FLOAT16IMAGE1开启 2D 图像类优化路径面向 GPU 纹理/图像内存布局FLOAT161让图像路径使用 float16 而非 float32二者通常配合使用以节省显存、提升带宽。ALLOW_TF32在 Ampere 及更新 NVIDIA GPU 上允许以 TensorFloat-32 精度执行矩阵乘换取吞吐提升代价是精度下降。BEAM取值 N 表示内核 beam search 时使用 N 条候选路径。beam search 发生在 AST 到 UOps 的 lowering 阶段源码入口位于 tinygrad/codegen/opt/postrange.py其中with Context(ALLOW_DEVICE_USAGE1): k beam_search(k, rawbufs, var_vals, beam, bool(getenv(BEAM_ESTIMATE, 1)))。BEAM_ESTIMATE默认 1表示先用估算快速淘汰劣质候选再对剩余候选做真实测量。WEBGPU_BACKEND 与 CUDA_PATH 的实现WEBGPU_BACKEND在 tinygrad/runtime/ops_webgpu.py 中被读取adapter_res InstanceRequestAdapter(instance, webgpu.WGPURequestAdapterOptions( powerPreferencewebgpu.WGPUPowerPreference_HighPerformance, backendTypebackend_types.get(getenv(WEBGPU_BACKEND, ), 0)))取值包括WGPUBackendType_Metal、WGPUBackendType_DirectX、WGPUBackendType_OpenGL、WGPUBackendType_Vulkan等用于在支持多后端的 WebGPU 环境中强制选择底层图形 API。CUDA_PATH在 tinygrad/runtime/support/compiler_cuda.py 中定义为CUDA_PATH getenv(CUDA_PATH, )用于定位CUDA_PATH/include下的 CUDA 头文件未设置时按文档所述回退到/usr/local/cuda/include、/usr/include与/opt/cuda/include。DEV 变量详解目标三元组与接口选择DEV是所有环境变量中语法最复杂的官方文档专门为其开辟了一节。它的完整语法是设备device: 渲染器renderer: 架构arch三段式段与段之间用冒号分隔device目标硬件平台如AMD、NV、CUDA、CPU、CL、METAL等。tinygrad 支持的后端全集定义在 tinygrad/device.pyALL_DEVICES [METAL, AMD, NV, CUDA, QCOM, CL, CPU, DSP, WEBGPU]renderer可选目标渲染器如LLVM、CUDA等省略时 tinygrad 自动挑选可用渲染器。arch可选目标架构如sm_70、gfx950省略时自动推断。此外还可以用加号在目标三元组之前指定访问设备的接口例如USBAMD表示通过 USB 接口访问 AMD 设备。接口与三元组均可省略省略部分由 tinygrad 自动确定。该解析逻辑在源码Target.parse中实现tinygrad/helpers.pystaticmethod def parse(s:str) - Target: if len(split:s.split()) 2: (iface, indices), s ((iface_split[0], iface_split[1]) if len(iface_split:split[0].rsplit(:, 1)) 2 else (split[0], ), split[1]) elif len(split) 2: raise RuntimeError(ftoo many in target string: {s!r}) else: iface, indices , match [x.upper() if i 2 else x for i,x in enumerate(s.split(:))]: case [dev, ren, arch]: return Target(dev, ren, arch, iface, indices) case [dev, ren]: return Target(dev, ren, interfaceiface, indicesindices) case [dev]: return Target(dev, interfaceiface, indicesindices) case _: raise RuntimeError(ftoo many : in target string: {s!r})Target是一个包含device、renderer、arch、interface、indices五个字段的 frozen dataclasstinygrad/helpers.py其中indices对应文档示例中CPU:LLVM:x86_64,znver2,avx2,-avx512f里的 CPU 特性标志列表。_DEV的 value setter 会把字符串按分号拆分为多个 Targettinygrad/helpers.py支持同时指定多个设备。DEV 取值示例与解释下表为官方文档原始示例逐行解读DEVcontentsInterpretationAMDuse the AMD deviceAMD:LLVMuse the AMD device with the LLVM rendererNV:CUDA:sm_70use the NV device with the CUDA renderer targetting sm_70AMD::gfx950use the AMD device targetting gfx950USBAMDuse the AMD device over the USB interfaceCPU:LLVMuse the CPU device with the LLVM rendererCPU:LLVM:x86_64,znver2,avx2,-avx512fuse the CPU device with the LLVM renderer, with additional arch flags注意AMD::gfx950中两个冒号之间留空表示跳过渲染器、只指定架构CPU:LLVM:x86_64,znver2,avx2,-avx512f中逗号分隔的 CPU 特性标志支持-前缀来禁用某特性如-avx512f详细的 CPU arch 标志说明见 运行时文档。DEV 在设备初始化链路中的角色DEV值在设备初始化时被消费Device.DEFAULT返回DEV.device or self._select_devicetinygrad/device.py即指定了DEV就优先使用它否则遍历ALL_DEVICES自动选择第一个可用设备。渲染器的选择通过DEV.target(...)完成tinygrad/device.py接口如 PCI、USB的选择同样依赖DEV.targettinygrad/device.py。因此DEV一份配置即可同时约束设备、渲染器、架构与访问接口四个维度。Debug breakdownDEBUG 各级输出的源码印证DEBUG是使用频率最高的调试开关其输出从 1 到 7 逐级递进。下表为官方文档完整内容结合源码说明每一级的实际输出位置VariableValueDescriptionDEBUG 1Enables debugging and lists devices being usedDEBUG 2Provides performance metrics for operations, including timing, memory usage, bandwidth for each kernel executionDEBUG 3Outputs the applied optimizations at a kernel levelDEBUG 4Outputs the generated kernel codeDEBUG 5Displays the intermediate representation of the computation UOpsDEBUG 6Displays the intermediate representation of the computation UOps in a linearized manner, detailing the operation sequenceDEBUG 7Outputs the assembly code generated for the target hardware各等级在源码中的落点DEBUG 1设备打开与编译进度提示。Device.__get_canonicalized_item中if DEBUG 1: print(fopened device {ix} from pid:{os.getpid()})tinygrad/device.py编译阶段tqdm(..., disableDEBUG1)控制进度条显隐tinygrad/engine/realize.py。DEBUG 2每个内核的性能指标。run_linear中if DEBUG 2 and not PROFILE: return直接跳过统计输出随后打印每次 kernel 执行的 timing、memory、bandwidth 明细tinygrad/engine/realize.py同时hcq_compile的 profile 开关为bool(PROFILE or DEBUG 2)tinygrad/engine/realize.py即 DEBUG2 时会顺带开启 HCQ 硬件队列的 profile 记录。DEBUG 4打印生成的内核代码。例如指令选择阶段结束后if DEBUG 4: print(ctx.asm_str(lst, sink.arg.function_name))tinygrad/codegen/init.py。DEBUG 7输出面向目标硬件的汇编代码同时在缓冲区分配/释放时打印内存操作if DEBUG is not None and DEBUG 7: print(fbuffer: deallocate {self.nbytes} bytes on {self.device})见 tinygrad/device.py。典型调试组合示例# 查看设备选择与每个内核的性能指标 DEVAMD DEBUG2 python3 examples/beautiful_mnist.py # 打印生成的 CUDA 内核源码 DEVNV:CUDA DEBUG4 python3 examples/beautiful_mnist.py # 输出中间表示 UOps线性化 DEBUG6 python3 -c from tinygrad import Tensor; (Tensor.ones(4,4)Tensor.ones(4,4)).realize() # 结合 beam search 观察内核优化过程 BEAM2 DEBUG3 python3 examples/beautiful_mnist.py综合实践把环境变量组合起来多设备并行DEV支持用分号指定多个目标_DEV.valuesetter 按;拆分例如DEVCPU;AMD可同时使用 CPU 与 AMD 设备配合多设备张量 API 使用。复现确定性调度调试调度问题时用JIT0关闭 JIT、用CCACHE0关闭编译器缓存tinygrad/helpers.py避免缓存干扰对真实编译路径的观察。性能剖析VIZ1隐含开启 PROFILE可配合 tinygrad 自带的 viz 可视化模块 查看内核执行时间线单独使用PROFILE1则只做采样不做可视化。内存监控DEBUG2 的带宽输出配合GlobalCounters定义于 tinygrad/helpers.py可观察全局内存占用与带宽变化。小结tinygrad 把运行时开关统一收敛为环境变量 ContextVar的双通道机制环境变量设定进程级默认值Context/with Context(...)提供函数级或代码块级的临时覆盖。掌握DEV的目标三元组语法、DEBUG的 1–7 级递进输出、以及BEAM、JIT、DEFAULT_FLOAT等行为开关就能在不改动一行代码的情况下完成设备切换、内核调试与性能分析。全部变量的权威定义与默认值可以在 tinygrad/helpers.py 中交叉验证相关架构背景可进一步阅读 开发者文档 与 运行时文档。【免费下载链接】tinygradYou like pytorch? You like micrograd? You love tinygrad! ❤️项目地址: https://gitcode.com/GitHub_Trending/tiny/tinygrad创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考