be faster to get the training result with triton lib

[JudeGu]

There is a optional argv in train_full_sft.py script of "--use_compile". If set this argv to True. The triton library will be needed to excelerate the training process almostly x1.5 speed.
I attach the whl I had search.
https://pypi.org/project/triton-windows/#files
https://pypi.org/project/triton/#files

I tested the 3.1.0, but the train_full_sft.py script raise the error below
After that I use the 3.8.0 .It running normally.

BR

(.venv) E:\backup\software\minimind\minimind\trainer>python train_full_sft.py
Model Params: 63.91M
Trainable Params: 63.912M
Loading dataset shards: 100%|██████████████████████████████████████████████████████████| 27/27 [00:00<00:00, 79.49it/s]
torch.compile enabled
Epoch [1/5]: 跳过前5000个step，从step 5001开始
main.c
正在创建库 C:\Users\admin\AppData\Local\Temp\tmpi8agysif\cuda_utils.cp311-win_amd64.lib 和对象 C:\Users\admin\AppData\Local\Temp\tmpi8agysif\cuda_utils.cp311-win_amd64.exp
E1008 23:04:02.638000 2776 Lib\site-packages\torch_inductor\runtime\triton_heuristics.py:1417] [1/0] Triton compilation failed: triton_poi_fused_embedding_0
E1008 23:04:02.638000 2776 Lib\site-packages\torch_inductor\runtime\triton_heuristics.py:1417] [1/0] def triton_poi_fused_embedding_0(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK : tl.constexpr):
E1008 23:04:02.638000 2776 Lib\site-packages\torch_inductor\runtime\triton_heuristics.py:1417] [1/0] xnumel = 9437184
E1008 23:04:02.638000 2776 Lib\site-packages\torch_inductor\runtime\triton_heuristics.py:1417] [1/0] xoffset = tl.program_id(0) * XBLOCK
E1008 23:04:02.638000 2776 Lib\site-packages\torch_inductor\runtime\triton_heuristics.py:1417] [1/0] xindex = xoffset + tl.arange(0, XBLOCK)[:]
E1008 23:04:02.638000 2776 Lib\site-packages\torch_inductor\runtime\triton_heuristics.py:1417] [1/0] xmask = tl.full([XBLOCK], True, tl.int1)[:]
E1008 23:04:02.638000 2776 Lib\site-packages\torch_inductor\runtime\triton_heuristics.py:1417] [1/0] x1 = xindex // 768
E1008 23:04:02.638000 2776 Lib\site-packages\torch_inductor\runtime\triton_heuristics.py:1417] [1/0] x0 = (xindex % 768)
E1008 23:04:02.638000 2776 Lib\site-packages\torch_inductor\runtime\triton_heuristics.py:1417] [1/0] x2 = xindex
E1008 23:04:02.638000 2776 Lib\site-packages\torch_inductor\runtime\triton_heuristics.py:1417] [1/0] tmp0 = tl.load(in_ptr0 + (x1), None, eviction_policy='evict_last')
E1008 23:04:02.638000 2776 Lib\site-packages\torch_inductor\runtime\triton_heuristics.py:1417] [1/0] tl.device_assert((0 <= tmp0) & (tmp0 < 6400), "index out of bounds: 0 <= tmp0 < 6400")
E1008 23:04:02.638000 2776 Lib\site-packages\torch_inductor\runtime\triton_heuristics.py:1417] [1/0] tmp2 = tl.load(in_ptr1 + (x0 + 768*tmp0), None)
E1008 23:04:02.638000 2776 Lib\site-packages\torch_inductor\runtime\triton_heuristics.py:1417] [1/0] tl.store(out_ptr0 + (x2), tmp2, None)
E1008 23:04:02.638000 2776 Lib\site-packages\torch_inductor\runtime\triton_heuristics.py:1417] [1/0]
E1008 23:04:02.638000 2776 Lib\site-packages\torch_inductor\runtime\triton_heuristics.py:1417] [1/0] metadata: {'signature': {'in_ptr0': '*i64', 'in_ptr1': '*fp32', 'out_ptr0': '*fp32', 'xnumel': 'i32'}, 'device': 0, 'constants': {'XBLOCK': 1024}, 'native_matmul': False, 'enable_fp_fusion': True, 'launch_pdl': False, 'disable_ftz': False, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1, 2, 3), equal_to_1=())], 'device_type': 'cuda', 'num_warps': 4, 'num_stages': 1, 'debug': True, 'cc': 120}
E1008 23:04:02.638000 2776 Lib\site-packages\torch_inductor\runtime\triton_heuristics.py:1417] [1/0] Traceback (most recent call last):
E1008 23:04:02.638000 2776 Lib\site-packages\torch_inductor\runtime\triton_heuristics.py:1417] [1/0] File "E:\backup\software\minimind\minimind.venv\Lib\site-packages\torch_inductor\runtime\triton_heuristics.py", line 1415, in _precompile_config
E1008 23:04:02.638000 2776 Lib\site-packages\torch_inductor\runtime\triton_heuristics.py:1417] [1/0] binary = triton.compile(*compile_args, **compile_kwargs)
E1008 23:04:02.638000 2776 Lib\site-packages\torch_inductor\runtime\triton_heuristics.py:1417] [1/0] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E1008 23:04:02.638000 2776 Lib\site-packages\torch_inductor\runtime\triton_heuristics.py:1417] [1/0] File "E:\backup\software\minimind\minimind.venv\Lib\site-packages\triton\compiler\compiler.py", line 286, in compile
E1008 23:04:02.638000 2776 Lib\site-packages\torch_inductor\runtime\triton_heuristics.py:1417] [1/0] next_module = compile_ir(module, metadata)
E1008 23:04:02.638000 2776 Lib\site-packages\torch_inductor\runtime\triton_heuristics.py:1417] [1/0] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E1008 23:04:02.638000 2776 Lib\site-packages\torch_inductor\runtime\triton_heuristics.py:1417] [1/0] File "E:\backup\software\minimind\minimind.venv\Lib\site-packages\triton\backends\nvidia\compiler.py", line 341, in
E1008 23:04:02.638000 2776 Lib\site-packages\torch_inductor\runtime\triton_heuristics.py:1417] [1/0] stages["ptx"] = lambda src, metadata: self.make_ptx(src, metadata, options, self.capability)
E1008 23:04:02.638000 2776 Lib\site-packages\torch_inductor\runtime\triton_heuristics.py:1417] [1/0] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E1008 23:04:02.638000 2776 Lib\site-packages\torch_inductor\runtime\triton_heuristics.py:1417] [1/0] File "E:\backup\software\minimind\minimind.venv\Lib\site-packages\triton\backends\nvidia\compiler.py", line 267, in make_ptx
E1008 23:04:02.638000 2776 Lib\site-packages\torch_inductor\runtime\triton_heuristics.py:1417] [1/0] ptx_version = ptx_get_version(cuda_version)
E1008 23:04:02.638000 2776 Lib\site-packages\torch_inductor\runtime\triton_heuristics.py:1417] [1/0] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E1008 23:04:02.638000 2776 Lib\site-packages\torch_inductor\runtime\triton_heuristics.py:1417] [1/0] File "E:\backup\software\minimind\minimind.venv\Lib\site-packages\triton\backends\nvidia\compiler.py", line 62, in ptx_get_version
E1008 23:04:02.638000 2776 Lib\site-packages\torch_inductor\runtime\triton_heuristics.py:1417] [1/0] raise RuntimeError("Triton only support CUDA 10.0 or higher, but got CUDA version: " + cuda_version)
E1008 23:04:02.638000 2776 Lib\site-packages\torch_inductor\runtime\triton_heuristics.py:1417] [1/0] RuntimeError: Triton only support CUDA 10.0 or higher, but got CUDA version: 13.4
Traceback (most recent call last):
File "E:\backup\software\minimind\minimind\trainer\train_full_sft.py", line 158, in
train_epoch(epoch, loader, len(loader) + skip, start_step, wandb)
File "E:\backup\software\minimind\minimind\trainer\train_full_sft.py", line 34, in train_epoch
res = model(input_ids, labels=labels)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "E:\backup\software\minimind\minimind.venv\Lib\site-packages\torch_dynamo\eval_frame.py", line 563, in call
return super().call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "E:\backup\software\minimind\minimind.venv\Lib\site-packages\torch\nn\modules\module.py", line 1783, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "E:\backup\software\minimind\minimind.venv\Lib\site-packages\torch\nn\modules\module.py", line 1794, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "E:\backup\software\minimind\minimind.venv\Lib\site-packages\torch_dynamo\eval_frame.py", line 1278, in compile_wrapper
raise e.remove_dynamo_frames() from None # see TORCHDYNAMO_VERBOSE=1
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "E:\backup\software\minimind\minimind.venv\Lib\site-packages\torch_inductor\compile_fx.py", line 1192, in _compile_fx_inner
raise InductorError(e, currentframe()).with_traceback(
File "E:\backup\software\minimind\minimind.venv\Lib\site-packages\torch_inductor\compile_fx.py", line 1169, in _compile_fx_inner
mb_compiled_graph = fx_codegen_and_compile(
^^^^^^^^^^^^^^^^^^^^^^^
File "E:\backup\software\minimind\minimind.venv\Lib\site-packages\torch_inductor\compile_fx.py", line 2001, in fx_codegen_and_compile
return scheme.codegen_and_compile(gm, example_inputs, inputs_to_check, graph_kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "E:\backup\software\minimind\minimind.venv\Lib\site-packages\torch_inductor\compile_fx.py", line 1749, in codegen_and_compile
compiled_module = graph.compile_to_module()
^^^^^^^^^^^^^^^^^^^^^^^^^
File "E:\backup\software\minimind\minimind.venv\Lib\site-packages\torch_inductor\graph.py", line 3060, in compile_to_module
return self._compile_to_module()
^^^^^^^^^^^^^^^^^^^^^^^^^
File "E:\backup\software\minimind\minimind.venv\Lib\site-packages\torch_inductor\graph.py", line 3070, in _compile_to_module
mod = self._compile_to_module_lines(wrapper_code)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "E:\backup\software\minimind\minimind.venv\Lib\site-packages\torch_inductor\graph.py", line 3146, in _compile_to_module_lines
mod = PyCodeCache.load_by_key_path(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "E:\backup\software\minimind\minimind.venv\Lib\site-packages\torch_inductor\codecache.py", line 4820, in load_by_key_path
mod = _reload_python_module(key, path, set_sys_modules=set_sys_modules)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "E:\backup\software\minimind\minimind.venv\Lib\site-packages\torch_inductor\runtime\compile_tasks.py", line 37, in _reload_python_module
exec(code, mod.dict, mod.dict)
File "C:\tc\ih\cihjtnefntojfhnvx6wvohhhcmp63cnuebw6xootvcsfv6cfy67p.py", line 48, in
triton_poi_fused_embedding_0 = async_compile.triton('triton_poi_fused_embedding_0', '''
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "E:\backup\software\minimind\minimind.venv\Lib\site-packages\torch_inductor\async_compile.py", line 584, in triton
kernel.precompile(
File "E:\backup\software\minimind\minimind.venv\Lib\site-packages\torch_inductor\runtime\triton_heuristics.py", line 804, in precompile
self._precompile_worker()
File "E:\backup\software\minimind\minimind.venv\Lib\site-packages\torch_inductor\runtime\triton_heuristics.py", line 827, in _precompile_worker
compile_results.append(self._precompile_config(c))
^^^^^^^^^^^^^^^^^^^^^^^^^^
File "E:\backup\software\minimind\minimind.venv\Lib\site-packages\torch_inductor\runtime\triton_heuristics.py", line 1415, in _precompile_config
binary = triton.compile(*compile_args, **compile_kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "E:\backup\software\minimind\minimind.venv\Lib\site-packages\triton\compiler\compiler.py", line 286, in compile
next_module = compile_ir(module, metadata)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "E:\backup\software\minimind\minimind.venv\Lib\site-packages\triton\backends\nvidia\compiler.py", line 341, in
stages["ptx"] = lambda src, metadata: self.make_ptx(src, metadata, options, self.capability)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "E:\backup\software\minimind\minimind.venv\Lib\site-packages\triton\backends\nvidia\compiler.py", line 267, in make_ptx
ptx_version = ptx_get_version(cuda_version)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "E:\backup\software\minimind\minimind.venv\Lib\site-packages\triton\backends\nvidia\compiler.py", line 62, in ptx_get_version
raise RuntimeError("Triton only support CUDA 10.0 or higher, but got CUDA version: " + cuda_version)
torch._inductor.exc.InductorError: RuntimeError: Triton only support CUDA 10.0 or higher, but got CUDA version: 13.4

Set TORCHDYNAMO_VERBOSE=1 for the internal stack trace (please do this especially if you're reporting a bug to PyTorch). For even more developer context, set TORCH_LOGS="+dynamo"