tinygrad 环境变量完全指南：DEV、DEBUG、BEAM 等运行时开关的用法与源码级解读
【免费下载链接】tinygradYou like pytorch? You like micrograd? You love tinygrad! ❤️项目地址: https://gitcode.com/GitHub_Trending/tiny/tinygrad
tinygrad 将绝大多数运行时行为开关设计为环境变量：从选择后端设备（DEV）、控制调试输出粒度（DEBUG）、调整内核搜索策略（BEAM）到切换默认浮点精度（DEFAULT_FLOAT），均可在不修改任何代码的情况下通过一行环境变量生效。本文以官方文档 docs/env_vars.md 为骨架，结合 tinygrad 源码（tinygrad/helpers.py、tinygrad/device.py、tinygrad/engine/realize.py 等）逐项讲解每个变量的取值、作用与底层实现，读完即可熟练使用DEV=AMD:LLVM DEBUG=4这类组合进行设备选择、性能分析和内核调试。
环境变量的两种生效方式
tinygrad 的环境变量既可以在进程启动前通过 shell 设置，也可以在代码运行期内临时切换。官方文档给出的进程级用法是：
DEV=CL DEBUG=4 python3 -m pytest即在命令行前缀中一次性设置多个变量，覆盖整个进程的运行时行为。
用@Context装饰器限定单个函数
对于 tinygrad 开发者，可以在函数上使用@Context(...)装饰器，让某个环境变量只在函数体内生效：
# in tensor.py (probably only useful if you are a tinygrad developer) @Context(DEBUG=4) 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(DEBUG=0): a = Tensor.ones(10, 10) a *= 2从源码看，所有可配置变量都是ContextVar的实例（定义于 tinygrad/helpers.py）。ContextVar.__init__会调用getenv(key, default_value)读取进程环境变量作为初始值，因此环境变量是全局默认值，而Context是局部覆盖，二者共享同一套变量名。
getenv的实现也值得一提（tinygrad/helpers.py）：
@functools.cache def getenv(key:str, default:Any=0): return type(default)(os.getenv(key, default))它会将字符串环境变量强制转换为默认值对应的类型（如int、str、float），并带functools.cache缓存，这解释了为什么文档表格中的整数值可以直接参与比较运算。
全局变量表：控制核心运行时行为的开关
以下变量控制 tinygrad 作为库被使用时（import tinygrad之后）的核心行为。表中#表示该变量可取任意整数值。表格内容完整继承自 docs/env_vars.md，并补充了源码中的默认值与实现位置。
| Variable | Possible Value(s) | Description |
|---|---|---|
| DEBUG | [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 search |
| DEFAULT_FLOAT | [HALF, ...] | specify the default float dtype (FLOAT32, HALF, BFLOAT16, FLOAT64, ...), default to FLOAT32 |
| IMAGE | [1] | enable 2d specific optimizations |
| FLOAT16 | [1] | use float16 for images instead of float32 |
| JIT | [0-2] | 0=disabled, 1=jit enabled (default), 2=jit enabled, but graphs are disabled |
| VIZ | [1] | 0=disabled, 1=viz enabled |
| ALLOW_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_PATH | str | UseCUDA_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", "")，默认空字符串，表示自动选择设备；DEBUG：ContextVar("DEBUG", 0)，默认关闭调试；BEAM：ContextVar("BEAM", 0)，默认不做 beam search；NOOPT：ContextVar("NOOPT", 0)。IMAGE、FLOAT16默认均为 0；JIT默认为 1（启用 JIT），JIT_BATCH_SIZE默认为 32。DEFAULT_FLOAT默认为"float32"，DEFAULT_INT默认为"int32"，二者由 dtype 系统在创建张量时读取。VIZ默认为 0；源码注释指出 "VIZ implies PROFILE, but you can run PROFILE without VIZ"，即PROFILE = ContextVar("PROFILE", abs(VIZ.value))会随 VIZ 自动开启性能采样。ALLOW_TF32默认为 0（tinygrad/helpers.py），源码注释明确其用途是 "allow tf32 to be used on NVIDIA GPUs"。- 文档未列出但同处定义的
CCACHE（默认 1，编译器缓存）与SCACHE（默认 1，调度器缓存）也可用相同方式在命令行覆盖。
几个易混淆变量的实际作用
- JIT（0/1/2）：默认 1 启用 JIT 编译，将多次 realize 合并为一次内核启动；设为 2 时 JIT 仍开启但禁用计算图（graph）记录，通常用于排查 graph 相关行为；设为 0 则完全关闭 JIT，每次 realize 立即执行。
- IMAGE 与 FLOAT16：
IMAGE=1开启 2D 图像类优化路径（面向 GPU 纹理/图像内存布局），FLOAT16=1让图像路径使用 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_USAGE=1): 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( powerPreference=webgpu.WGPUPowerPreference_HighPerformance, backendType=backend_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.py：ALL_DEVICES = ["METAL", "AMD", "NV", "CUDA", "QCOM", "CL", "CPU", "DSP", "WEBGPU"]renderer（可选）：目标渲染器，如
LLVM、CUDA等；省略时 tinygrad 自动挑选可用渲染器。arch（可选）：目标架构，如
sm_70、gfx950；省略时自动推断。
此外还可以用加号（+）在目标三元组之前指定访问设备的接口，例如USB+AMD表示通过 USB 接口访问 AMD 设备。接口与三元组均可省略，省略部分由 tinygrad 自动确定。
该解析逻辑在源码Target.parse中实现（tinygrad/helpers.py）：
@staticmethod 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(f"too 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, interface=iface, indices=indices) case [dev]: return Target(dev, interface=iface, indices=indices) case _: raise RuntimeError(f"too many ':' in target string: {s!r}")Target是一个包含device、renderer、arch、interface、indices五个字段的 frozen dataclass（tinygrad/helpers.py），其中indices对应文档示例中CPU:LLVM:x86_64,znver2,avx2,-avx512f里的 CPU 特性标志列表。_DEV的 value setter 会把字符串按分号拆分为多个 Target（tinygrad/helpers.py），支持同时指定多个设备。
DEV 取值示例与解释
下表为官方文档原始示例，逐行解读：
DEVcontents | Interpretation |
|---|---|
| AMD | use the AMD device |
| AMD:LLVM | use the AMD device with the LLVM renderer |
| NV:CUDA:sm_70 | use the NV device with the CUDA renderer targetting sm_70 |
| AMD::gfx950 | use the AMD device targetting gfx950 |
| USB+AMD | use the AMD device over the USB interface |
| CPU:LLVM | use the CPU device with the LLVM renderer |
| CPU:LLVM:x86_64,znver2,avx2,-avx512f | use 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_device（tinygrad/device.py），即指定了DEV就优先使用它，否则遍历ALL_DEVICES自动选择第一个可用设备。渲染器的选择通过DEV.target(...)完成（tinygrad/device.py），接口（如 PCI、USB）的选择同样依赖DEV.target（tinygrad/device.py）。因此DEV一份配置即可同时约束设备、渲染器、架构与访问接口四个维度。
Debug breakdown：DEBUG 各级输出的源码印证
DEBUG是使用频率最高的调试开关，其输出从 1 到 7 逐级递进。下表为官方文档完整内容，结合源码说明每一级的实际输出位置：
| Variable | Value | Description |
|---|---|---|
| DEBUG | >= 1 | Enables debugging and lists devices being used |
| DEBUG | >= 2 | Provides performance metrics for operations, including timing, memory usage, bandwidth for each kernel execution |
| DEBUG | >= 3 | Outputs the applied optimizations at a kernel level |
| DEBUG | >= 4 | Outputs the generated kernel code |
| DEBUG | >= 5 | Displays the intermediate representation of the computation UOps |
| DEBUG | >= 6 | Displays the intermediate representation of the computation UOps in a linearized manner, detailing the operation sequence |
| DEBUG | >= 7 | Outputs the assembly code generated for the target hardware |
各等级在源码中的落点：
- DEBUG >= 1：设备打开与编译进度提示。
Device.__get_canonicalized_item中if DEBUG >= 1: print(f"opened device {ix} from pid:{os.getpid()}")（tinygrad/device.py）；编译阶段tqdm(..., disable=DEBUG<1)控制进度条显隐（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），即 DEBUG>=2 时会顺带开启 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(f"buffer: deallocate {self.nbytes} bytes on {self.device}")，见 tinygrad/device.py）。
典型调试组合示例：
# 查看设备选择与每个内核的性能指标 DEV=AMD DEBUG=2 python3 examples/beautiful_mnist.py # 打印生成的 CUDA 内核源码 DEV=NV:CUDA DEBUG=4 python3 examples/beautiful_mnist.py # 输出中间表示 UOps（线性化） DEBUG=6 python3 -c "from tinygrad import Tensor; (Tensor.ones(4,4)@Tensor.ones(4,4)).realize()" # 结合 beam search 观察内核优化过程 BEAM=2 DEBUG=3 python3 examples/beautiful_mnist.py综合实践：把环境变量组合起来
- 多设备并行：
DEV支持用分号指定多个目标（_DEV.valuesetter 按;拆分），例如DEV=CPU;AMD可同时使用 CPU 与 AMD 设备，配合多设备张量 API 使用。 - 复现确定性调度：调试调度问题时用
JIT=0关闭 JIT、用CCACHE=0关闭编译器缓存（tinygrad/helpers.py），避免缓存干扰对真实编译路径的观察。 - 性能剖析：
VIZ=1隐含开启 PROFILE，可配合 tinygrad 自带的 viz 可视化模块 查看内核执行时间线；单独使用PROFILE=1则只做采样不做可视化。 - 内存监控：DEBUG>=2 的带宽输出配合
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），仅供参考
