在 SCNet 海光 BW 64GB 实例上跑通 PyTorch、Triton 与 FlashAttention：环境检查与实测记录

https://www.scnet.cn/ui/console/index.html#/notebook/add?resourceGroupNamehx1hgbwnormalclusterId20091拿到这台海光 BW 实例后我没有马上安装框架或下载模型而是先确认三件事设备是否真的被计算框架识别、容器实际能用多少 CPU/内存以及现有软件栈能不能完成前向和反向计算。这一步很重要。平台规格、容器里lscpu/free看到的宿主机资源以及作业受 cgroup 限制后真正可用的资源并不一定是同一回事。如果一开始把资源边界判断错了后面出现 DataLoader 卡顿、编译进程被杀或训练 OOM很容易误判成加速卡或框架问题。本文记录一次实际排查和验证过程。所有结果来自同一台 SCNet Notebook 实例截图做过脱敏文中的吞吐数据只用于验证计算链路是否正常不代表产品峰值性能。一、先说结论这套环境可以胜任以下工作PyTorch 基础教学和深度学习课程作业Transformer、RoPE、RMSNorm、SwiGLU、Attention 等算子的实现与调试BF16/FP16 混合精度实验Triton 自定义算子练习FlashAttention 前向、反向验证中小模型推理以及在显存允许范围内进行 LoRA/QLoRA 微调使用 vLLM 做单卡推理服务实验。本次实例的核心环境如下。项目实测结果操作系统Ubuntu 22.04.5 LTS加速卡海光 BW1 卡显存63.98 GiBDTK/opt/dtk-26.04-DCC2602-0317Python3.10.12PyTorch2.10.0HIP 运行时6.3.26093Triton3.5.0Triton 目标hip / gfx936 / warp_size64FlashAttention2.8.3 定制版本vLLM0.18.1BF16支持二、连接实例浏览器适合试用VS Code 更适合持续开发SCNet 页面提供 JupyterLab、工具面板和 SSH 登录信息。临时运行几条命令时浏览器终端已经够用如果要反复修改训练脚本、查看 Git diff 或调试 PythonVS Code Remote - SSH 会顺手很多。本地~/.ssh/config可以写成下面这样主机和端口用控制台显示的值替换。密码不要写入配置文件也不要提交到仓库。Host scnet-hygon-bw HostName SSH_HOST User root Port SSH_PORT ServerAliveInterval 30 ServerAliveCountMax 3如果远端不能直接访问 VS Code 的下载站点可以在本地 VS Code 设置中加入{remote.SSH.localServerDownload:always,remote.SSH.useExecServer:false}这样会先在本机下载 VS Code Server再通过 SSH 传到远端能绕开远端网络访问受限的问题。连接成功后我习惯先打开一个单独的实验目录而不是直接在/root下堆脚本mkdir-p/root/hygon-bw-guide-democd/root/hygon-bw-guide-demo三、不要只看lscpu和free先检查系统和 cgroup 限额cat/etc/os-release lscpufree-h# cgroup v2cat/sys/fs/cgroup/cpu.maxcat/sys/fs/cgroup/memory.max这台容器里能看到宿主机约 128 个 CPU 线程和 503 GiB 内存但作业实际受到约 15 个 CPU 核、59 GiB 内存的限制。因此num_workers不应照着 128 个线程设置编译并行度也不宜直接拉满。我会再用下面的方式把限制换算成人能直接读懂的数字python -PY from pathlib import Path quota, period Path(/sys/fs/cgroup/cpu.max).read_text().split() memory int(Path(/sys/fs/cgroup/memory.max).read_text()) cpu_limit unlimited if quota max else int(quota) / int(period) print(fCPU limit: {cpu_limit}) print(fMemory limit: {memory / 1024**3:.2f} GiB) PY这个差异会直接影响三类参数DataLoader 的num_workers源码编译时的并行任务数CPU offload、数据缓存和内存映射文件的规模。四、DTK 环境必须在当前 Shell 中激活这台实例已经预装 DTK但新开的 Shell 不一定自动带上所需环境变量。先执行source/opt/dtk/env.sh然后检查 PyTorch 是否能看到设备python -PY import torch print(torch:, torch.__version__) print(hip:, torch.version.hip) print(available:, torch.cuda.is_available()) print(device_count:, torch.cuda.device_count()) print(device:, torch.cuda.get_device_name(0)) print(bf16:, torch.cuda.is_bf16_supported()) PY这里有一个容易困惑的地方在 DTK/HIP 软件栈中PyTorch 仍然沿用torch.cuda这套接口名。代码里出现devicecuda并不表示程序偷偷切到了 NVIDIA CUDA它只是 PyTorch 兼容接口的历史命名实际后端可以从torch.version.hip和设备名称确认。设备状态可以用hy-smi查看hy-smi如果输出开头出现rocm_smi_lib回退警告但后续仍能正常列出 HCU、温度、功耗和显存占用我会先把它当作管理库路径提示而不是立刻认定设备不可用。最终还要以 PyTorch 分配张量和实际计算是否成功为准。五、一份可复现的四段式验证脚本只打印torch.cuda.is_available() True还不够。我更关心四条链路FP16/BF16 矩阵乘是否能执行并且结果有限PyTorch 原生 SDPA 能否完成 BF16 因果注意力的前向和反向Triton 能否针对gfx936编译并运行自定义 kernelFlashAttention 能否完成 BF16 因果注意力的前向和反向。下面的脚本会把完整结果保存成 JSON方便后续比较镜像或依赖版本。#!/usr/bin/env python3from__future__importannotationsimportjsonimportplatformimporttimefrompathlibimportPathimporttorchimporttorch.nn.functionalasF RESULTS_PATHPath(__file__).with_name(validation_results.json)defgib(num_bytes:int)-float:returnnum_bytes/1024**3deftimed_matmul(dtype:torch.dtype,size:int2048,repeats:int8,)-dict:torch.manual_seed(336)lhstorch.randn((size,size),devicecuda,dtypedtype)rhstorch.randn((size,size),devicecuda,dtypedtype)for_inrange(2):_lhs rhs torch.cuda.synchronize()startedtime.perf_counter()for_inrange(repeats):resultlhs rhs torch.cuda.synchronize()elapsed(time.perf_counter()-started)/repeats tflops2*size**3/elapsed/1e12asserttorch.isfinite(result).all()return{dtype:str(dtype).removeprefix(torch.),shape:f{size}x{size},average_ms:round(elapsed*1000,3),tflops:round(tflops,3),finite:True,}defcheck_sdpa()-dict:torch.manual_seed(336)qtorch.randn((2,8,128,64),devicecuda,dtypetorch.bfloat16,requires_gradTrue,)ktorch.randn_like(q,requires_gradTrue)vtorch.randn_like(q,requires_gradTrue)outputF.scaled_dot_product_attention(q,k,v,is_causalTrue,)output.float().square().mean().backward()tensors(output,q.grad,k.grad,v.grad)assertall(bool(torch.isfinite(item).all())foritemintensors)return{shape:list(output.shape),dtype:bfloat16,forward_finite:True,backward_finite:True,}defcheck_triton()-dict:importtritonimporttriton.languageastltriton.jitdefadd_kernel(x_ptr,y_ptr,out_ptr,n_elements:tl.constexpr,BLOCK:tl.constexpr,):offsetstl.program_id(0)*BLOCKtl.arange(0,BLOCK)maskoffsetsn_elements xtl.load(x_ptroffsets,maskmask)ytl.load(y_ptroffsets,maskmask)tl.store(out_ptroffsets,xy,maskmask)torch.manual_seed(336)n_elements65_537xtorch.randn(n_elements,devicecuda,dtypetorch.float32)ytorch.randn_like(x)outputtorch.empty_like(x)grid(triton.cdiv(n_elements,256),)add_kernel[grid](x,y,output,n_elements,BLOCK256)torch.cuda.synchronize()max_errorfloat((output-(xy)).abs().max())asserttorch.allclose(output,xy)targettriton.runtime.driver.active.get_current_target()return{version:triton.__version__,target:str(target),elements:n_elements,max_abs_error:max_error,passed:True,}defcheck_flash_attention()-dict:fromflash_attnimportflash_attn_func torch.manual_seed(336)shape(2,128,8,64)qtorch.randn(shape,devicecuda,dtypetorch.bfloat16,requires_gradTrue,)ktorch.randn_like(q,requires_gradTrue)vtorch.randn_like(q,requires_gradTrue)outputflash_attn_func(q,k,v,causalTrue)output.float().square().mean().backward()asserttorch.isfinite(output).all()asserttorch.isfinite(q.grad).all()return{shape:list(output.shape),dtype:bfloat16,forward_finite:True,backward_finite:True,}defmain()-None:asserttorch.cuda.is_available(),DTK accelerator is not visibledevicetorch.cuda.get_device_properties(0)report{system:{python:platform.python_version(),torch:torch.__version__,hip:torch.version.hip,device:torch.cuda.get_device_name(0),device_count:torch.cuda.device_count(),memory_gib:round(gib(device.total_memory),2),bf16_supported:torch.cuda.is_bf16_supported(),},matmul:[timed_matmul(torch.float16),timed_matmul(torch.bfloat16),],sdpa:check_sdpa(),triton:check_triton(),flash_attention:check_flash_attention(),}RESULTS_PATH.write_text(json.dumps(report,ensure_asciiFalse,indent2),encodingutf-8,)print(json.dumps(report,ensure_asciiFalse,indent2))if__name____main__:main()运行方式source/opt/dtk/env.sh python hygon_bw_validation.py六、实测结果与正确解读方式本次运行的关键结果如下Device: BW | VRAM: 63.98 GiB | BF16: True Python: 3.10.12 | PyTorch: 2.10.0 | HIP: 6.3.26093 MatMul float16: 2048x2048 | 0.074 ms | 231.486 TFLOPS | PASS MatMul bfloat16: 2048x2048 | 0.076 ms | 225.437 TFLOPS | PASS SDPA BF16: output[2, 8, 128, 64] | forward/backward finite | PASS Triton: GPUTarget(backendhip, archgfx936, warp_size64) vector add | max error0.00e00 | PASS FlashAttention BF16: causal forward/backward finite | PASS这里最值得强调的是231.486 TFLOPS和225.437 TFLOPS来自规模为2048 × 2048、预热 2 次、重复 8 次的短时烟测。它没有固定频率和功耗也没有覆盖更大矩阵、长时间稳态、并发任务和不同转置布局因此不能当成官方峰值更不能直接拿来与其他平台做横向排名。这组数据真正能说明的是FP16 和 BF16 计算链路已经打通输出没有 NaN 或 InfPyTorch 原生 Attention 的反向传播正常Triton 能识别gfx936自定义向量加法误差为 0当前定制版 FlashAttention 可以完成 BF16 causal 前向和反向。如果要做严肃的性能报告还需要固定软件版本、时钟/功耗状态、矩阵规模、预热次数和统计口径并至少报告中位数、P95、显存峰值及稳态功耗。七、这台 64GB 单卡实例适合练什么1. 课程作业和算子实验实现线性层、Embedding、RMSNorm、RoPE、SwiGLU、Multi-Head Attention、交叉熵和优化器都没有问题。建议每个模块先写一个小规模 CPU 参考实现再与设备版本比较torch.testing.assert_close(actual.float().cpu(),expected.float().cpu(),rtol1e-2,atol1e-2,)这种做法比一上来跑完整训练更容易定位错误尤其适合检查 mask 方向、张量布局、广播和 dtype 转换。2. 小模型预训练可以用它完成小型语言模型从分词、数据加载、训练到采样的完整闭环。为了让实验稳定我会优先采用BF16 参数或自动混合精度梯度累积换取更大的有效 batchgradient checkpointing 降低激活占用训练前先用几十个 step 验证 loss 是否下降每个阶段记录峰值显存和 tokens/s。3. LoRA/QLoRA 与单卡推理64GB 显存给单卡实验留出了不错的空间但“能加载”不等于“能全参数训练”。以 7B 模型为例Adam 类优化器的参数、梯度和状态通常远超仅保存一份权重的空间全参数训练还要给激活和临时缓冲区留余量。LoRA/QLoRA 更符合这类单卡实例的使用方式具体上限仍取决于序列长度、batch size、量化算子和内核兼容性。当前环境已经有 vLLM可以用于验证模型加载、生成和服务接口。正式跑模型前建议先用本地已有的小模型或离线上传的权重做 smoke test不要一开始就搬几十 GB 的权重。4. Triton 内核练习本次测试显示 Triton 的后端目标为GPUTarget(backendhip, archgfx936, warp_size64)这意味着可以练习向量加、Softmax、LayerNorm、矩阵乘等内核。不过从其他平台迁移 kernel 时不要默认 warp size 一定是 32也不要照搬只针对某一架构调出的num_warps和 block size。先保证正确再结合实际目标做调优。八、离线环境中最容易踩的几个坑1. 外网域名能解析不代表 HTTPS 能访问本次环境中 DNS 可以工作但访问 PyPI、GitHub 和 Hugging Face 的 HTTPS 请求受限。因此直接执行pip install或git clone可能长时间等待后失败。更稳妥的流程是在可联网机器下载 wheel、源码包和模型文件记录文件哈希通过平台文件管理、JupyterLab 或scp上传使用本地路径安装把依赖版本和验证结果一起归档。例如# 联网机器pip download-dwheels-rrequirements.txt# 上传后在远端离线安装pipinstall--no-index --find-links ./wheels-rrequirements.txt2. 每个新终端都要确认 DTK 环境如果torch.cuda.is_available()突然从True变成False先检查是否忘记执行source/opt/dtk/env.sh不要急着重装 PyTorch。重装很可能破坏平台已经匹配好的定制依赖。3. 自定义算子包必须做前向和反向验证能import flash_attn只能证明 Python 找到了包不能证明内核与当前设备完全兼容。至少应覆盖目标 dtypecausal/non-causal 模式真实训练会用到的 head dimension前向输出有限反向梯度有限。4. 先保存 JSON再截图截图适合写文档但不适合自动比较。我把系统版本、矩阵乘耗时、目标架构和误差同时写入validation_results.json。以后更换镜像或升级 DTK 时重新运行同一脚本即可做差异比较。九、我的建议先建立最小闭环再扩大实验在异构加速环境中最省时间的顺序通常是确认 cgroup 资源边界激活厂商软件栈用小张量完成一次设备计算验证核心算子的前向和反向保存结构化结果再接入数据集、模型和分布式组件。这次验证之后我对这台实例的定位比较清楚它很适合深度学习教学、单卡工程验证、Triton 算子练习和中小规模模型实验如果目标是长时间大模型预训练或严格的跨平台性能对比则还需要更完整的监控、数据集、集群和基准测试方案。对我来说这次排查最有价值的并不是跑出一个好看的 TFLOPS 数字而是确认了从 DTK、PyTorch、Triton 到 FlashAttention 的每一层都能给出可复现、可解释的结果。环境是否真正可用最终还是要由一段能重复运行的代码来回答。