-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathfeed.xml
More file actions
11806 lines (11498 loc) · 977 KB
/
feed.xml
File metadata and controls
11806 lines (11498 loc) · 977 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
<?xml version="1.0" encoding="utf-8"?>
<rss version="2.0" xmlns:dc="http://purl.org/dc/elements/1.1/" xmlns:content="http://purl.org/rss/1.0/modules/content/" xmlns:atom="http://www.w3.org/2005/Atom">
<channel>
<title>Peihao Yang</title>
<link>https://forsworns.github.io//</link>
<description>Personal Blog</description>
<lastBuildDate>Sun, 08 Feb 2026 03:13:59 GMT</lastBuildDate>
<docs>https://validator.w3.org/feed/docs/rss2.html</docs>
<generator>https://github.com/jpmonette/feed</generator>
<image>
<title>Peihao Yang</title>
<url>https://forsworns.github.io/assets/logo.png</url>
<link>https://forsworns.github.io//</link>
</image>
<copyright>MIT License</copyright>
<atom:link href="https://forsworns.github.io//feed.xml" rel="self" type="application/rss+xml"/>
<item>
<title><![CDATA[HAMI 源码阅读]]></title>
<link>https://forsworns.github.io///zh/blogs/20241020/</link>
<guid>https://forsworns.github.io///zh/blogs/20241020/</guid>
<pubDate>Sun, 20 Oct 2024 00:00:00 GMT</pubDate>
<description><![CDATA[第四范式开源的通用 GPU 虚拟化组件]]></description>
<content:encoded><![CDATA[<p>[[toc]]</p>
<p>第四范式开源的通用 GPU 虚拟化组件,支持多家 GPU 产品,目前只做了切分功能,已进 CNCF。比较像之前腾讯开源的 <a href="https://github.com/tkestack/gpu-manager">https://github.com/tkestack/gpu-manager</a> 和对应的 <a href="https://github.com/tkestack/vcuda-controller%EF%BC%8C%E4%B9%8B%E5%89%8D%E4%B8%80%E7%9B%B4%E6%B2%A1%E4%BB%94%E7%BB%86%E8%AF%BB%EF%BC%8C%E5%88%9A%E5%A5%BD%E8%AF%BB%E4%B8%8B%E8%BF%99%E4%B8%AA%E8%BF%9B%E4%BA%86">https://github.com/tkestack/vcuda-controller,之前一直没仔细读,刚好读下这个进了</a> CNCF 的吧。</p>
<h1>HAMi-core</h1>
<p><a href="https://github.com/Project-HAMi/HAMi-core">https://github.com/Project-HAMi/HAMi-core</a><br>
基于主线 6b2aed490910db1a33c6575ba81b1ecd96fce5f4</p>
<h2>src/libvgpu.c</h2>
<p>劫持 <a href="http://libcuda.so">libcuda.so</a> 和 <a href="http://libnvidia-ml.so">libnvidia-ml.so</a> 的方法是通过劫持 dlsym 函数,如果用户查询的某个符号是 HAMI-core 可以拦截的,就返回对应的拦截函数。dlsym 则是通过暴露一个同名符号,通过 LD_PRELOAD 等方式对 libdl 做覆盖拦截。</p>
<h2>src/nvml</h2>
<h3>hook.c</h3>
<p>做了改动的一些 NVML API,查初始化的时候构造的真实函数指针表调用过去。最重要的就是 <code>_nvmlDeviceGetMemoryInfo</code> 这个实现。</p>
<p><code>_nvmlDeviceGetMemoryInfo</code> 里面 <code>nvmlDeviceGetIndex</code> 获取到设备的 id,通过 <code>nvml_to_cuda_map()</code> 转换成 <code>shared_region_info_t</code> 中用于标识虚拟设备的 id。根据 id 查询 <code>shared_region_info_t</code> 类型的全局单例 <code>region_info</code>,获取当前虚拟设备的显存限制和使用情况。</p>
<h3>nvml_entry.c</h3>
<p>没有做改动直接调用下去的 NVML API。</p>
<h2>src/multiprocess</h2>
<h3>shrreg_tool.c</h3>
<p>一个命令行小工具,支持几个选项:</p>
<ul>
<li>create_new:创建了一个文件 <code>/tmp/cudevshr.cache</code>,后面会被用来做跨进程的共享区域,它只是保证这个文件存在。</li>
<li>Suspend/resume:对所有运行中的被监控到的进程执行 <code>SIGUSR1</code> 和 <code>SIGUSR2</code> 分别用于恢复和挂起这些任务。</li>
</ul>
<h3>multiprocess_memory_limit.c</h3>
<p>这个文件里面比较杂,主要是 HAMI 的多进程资源使用情况的共享内存文件缓存、基于这个共享内存实现的显存管理、还有一些工具函数如host/container pid 转换、共享内存的加锁(lock_shrreg、unlock_shrreg<br>
),虽然看文件名只是做显存限制的。vcuda-controller 里面实现的比较简单,就是在一个文件里面存了下每个进程的 pid,然后每个 API 调用都会调用 <code>nvmlDeviceGetComputeRunningProcesses</code> 去查然后去对 pid 做匹配,为了省时,搜索的时候对 pid 做了二分,总体上开销还是比较高的。HAMI 这里则是通过直接创建一个多进程共享的资源消耗统计文件,进行了缓存,减少 NVML API 调用次数。这个共享文件会被 mmap 到每个进程内,也就是 <code>shared_region_t</code> 类型的 <code>region_info.shared_region</code>。</p>
<p>初始化过程中主要做了两件事 <code>try_create_shrreg()</code> 和 <code>init_proc_slot_withlock()</code>。<code>try_create_shrreg</code> 就是创建、初始化上面提到的共享文件的过程。<code>init_proc_slot_withlock</code> 则是对共享内存中当前进程的 slot 做了初始化。为 <code>SIGUSR1</code> 和 <code>SIGUSR1</code> 分别注册信号处理函数 <code>sig_restore_stub</code> 和 <code>sig_swap_stub</code>,用来恢复/暂停显存分配。</p>
<p>那么有了上面的共享内存,HAMI 中就不是通过 NVML 去查询显存占用了,而是通过 <code>get_gpu_memory_usage</code> 直接查询共享内存缓存中的统计数据。那么它又是如何收集这个统计数据的呢?这就涉及到了另一个比较有趣的函数是 <code>add_gpu_device_memory_usage</code>,它在涉及到显存剧烈变化的 CUDA API 执行时被调用,用来修改共享内存中的显存消耗统计数据,同时它还支持对显存进行分类,区分了是 CUcontext 相关、CUmodule相关、数据相关三类。但是读完代码发现,其实只有通过 cuda API 分配的数据显存被准确统计到了,不知道是不是他们内部实现没有开源出来。CUmodule 实际上完全没统计,CUcontext 则是用了一下初始化后 primary context 的消耗去计算,不过除了 primary context,一般上层框架也基本没有自己去创建 CUcontext 的(其他厂商,AMD ROCm 甚至实际上没有区分 device 和 context)。</p>
<h3>multiprocess_utilization_watcher.c</h3>
<p>对 cuda core 进行分配,与 vcuda-controller 中类似。</p>
<p><code>cuda_to_nvml_map</code> 变量定义在这里,在别的地方 extern 引用了。全局变量 <code>g_cycle</code> 为 10ms,<code>g_wait</code> 为 120ms。</p>
<p>几个重要的函数,<code>setspec()</code> 用于计算 cuda core 的总数。这里的 FACTOR 是 32,没太搞懂为啥,A100 每个 SM 是 64 个 cuda core,L20、H100 GPU 每个 SM 上是 128 个 cuda core,当然也可能这里表示的是 CUDA 一个 SM 上的 wrap size 是 32。over-subscription 的话有允许调度多个线程,所以是乘起来。感觉只是一个软性的估计,下面的限流器实现其实也可以看出来是允许超过限制的。</p>
<pre><code>int setspec() {
CHECK_CU_RESULT(cuDeviceGetAttribute(&g_sm_num,CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT,0));
CHECK_CU_RESULT(cuDeviceGetAttribute(&g_max_thread_per_sm,CU_DEVICE_ATTRIBUTE_MAX_THREADS_PER_MULTIPROCESSOR,0));
g_total_cuda_cores = g_max_thread_per_sm * g_sm_num * FACTOR;
return 0;
}
</code></pre>
<p><code>rate_limiter()</code> 是对 utilization 做限制的核心实现(注意 nvidia-smi 看到 utilization 是一个时分的统计数据,这里从 cuda core 的角度去限制实际上是空分的),在 <code>cuLaunchKernel</code> 被调用的时候先触发它,用来限流。加载核函数的 grid 参数被用来计算需要占用的 cuda core 数量。这里在限流的时候,留了一个优先级的接口,<code>region_info.shared_region->priority</code> 目前看没啥实际用途。(NSDI 23 的一个类似的工作 <a href="https://github.com/pkusys/TGS%EF%BC%8C%E5%81%9A%E4%BA%86%E4%BC%98%E5%85%88%E7%BA%A7%E8%B0%83%E5%BA%A6%EF%BC%8C%E7%AE%80%E5%8D%95%E7%9C%8B%E4%BA%86%E4%B8%8B%E4%BB%A3%E7%A0%81%E9%83%BD%E6%98%AF%E5%90%8C%E6%A0%B9%E7%94%9F~%EF%BC%89%E3%80%82%E7%B2%BE%E7%AE%80%E5%90%8E%E6%A0%B8%E5%BF%83%E5%A6%82%E4%B8%8B">https://github.com/pkusys/TGS,做了优先级调度,简单看了下代码都是同根生~)。精简后核心如下</a></p>
<pre><code class="language-cpp">// 同样没有用 block,按 vcuda-controller 那边的解释,他们是认为 block 的影响没有 grid 大所以只用了 grid,所以只是留了一个参数给其他算法实现
void rate_limiter(int grids, int blocks) {
int before_cuda_cores = 0;
int after_cuda_cores = 0;
int kernel_size = grids;
do {
before_cuda_cores = g_cur_cuda_cores;
// cuda core 不够了,进入睡眠,每 10ms 检查一次 cuda core 资源,等待别的核函数跑完释放了 cuda core,再提交新的核函数。
if (before_cuda_cores < 0) {
nanosleep(&g_cycle, NULL);
continue;
}
// 更新如果加载了核函数,会剩余的 cuda core 数量。按这种实现,用户可以使用远超过限制的 cuda core?
after_cuda_cores = before_cuda_cores - kernel_size;
} while (!CAS(&g_cur_cuda_cores, before_cuda_cores, after_cuda_cores));
}
</code></pre>
<p><code>utilization_watcher</code> 是一个进程内的守护线程,在 <code>src/libvgpu.c</code> 中通过 <code>init_utilization_watcher()</code> 在初始化完成后创建,以 <code>g_wait</code> 为周期重新分配 cuda core 计算资源,较 vcuda-controller 中的 cuda core 分配策略做了很大的删减,代码简化后如下,</p>
<pre><code class="language-cpp">void* utilization_watcher() {
int userutil[CUDA_DEVICE_MAX_COUNT];
int sysprocnum;
int share = 0;
int upper_limit = get_current_device_sm_limit(0);
while (1){
// 120ms 更新一次分配
nanosleep(&g_wait, NULL);
if (pidfound==0) {
// 在 `region_info.shared_region->procs` 中注册自己,目前代码中写死了最多支持 1024 个进程。
update_host_pid();
}
// 设置进程 sm 利用率为 0
init_gpu_device_sm_utilization();
// 这里实际上拿到了多卡的信息,而且和 vcuda-controller 不同的是它是会把查询结果写到一个共享文件里面做缓存。
get_used_gpu_utilization(userutil,&sysprocnum);
if ((share==g_total_cuda_cores) && (g_cur_cuda_cores<0)) {
// 这里没看懂
g_total_cuda_cores *= 2;
share = g_total_cuda_cores;
}
// 但是按这里的写法,当前是只支持了单卡,没有利用到上一步取出的多卡信息,
// 根据利用率限制、当前的利用率、上次的变化值,计算这次分配额度的变化
share = delta(upper_limit, userutil[0], share);
// 应用计算出的额度变化,重新配置 `g_cur_cuda_cores`
change_token(share);
}
}
</code></pre>
<p>上面通过在 <code>get_used_gpu_utilization()</code> 内,调用 <code>nvmlDeviceGetComputeRunningProcesses</code> 获取所有正在使用某个 GPU 设备的进程,然后调用 <code>nvmlDeviceGetProcessUtilization</code> 获取每个进程的 GPU 利用率和显存占用,分别记录到了 <code>region_info.shared_region->procs[i].device_util[dev].sm_util</code> 和 <code>region_info.shared_region->procs[i].monitorused[dev]</code>。这个函数是考虑了多进程的。</p>
<h2>src/allocator/allocator.c</h2>
<p><code>allocated_list</code> 是一个存了 <code>allocated_device_memory_struct</code> 的双向链表,实例化了两个全局变量 <code>device_overallocated</code>、<code>array_list</code>。成员定义如下</p>
<pre><code class="language-cpp">struct allocated_device_memory_struct{
CUdeviceptr address;
size_t length;
CUcontext ctx;
CUmemGenericAllocationHandle *allocHandle;
};
</code></pre>
<p><code>region_list</code> 是一个存了 <code>region_struct</code> 的双向链表,实例化成了一个全局变量 ``r_list。成员定义如下</p>
<pre><code class="language-cpp">struct region_struct{
size_t start;
size_t freemark;
size_t freed_map;
size_t length;
CUcontext ctx;
allocated_list *region_allocs;
char *bitmap;
CUmemGenericAllocationHandle *allocHandle;
};
</code></pre>
<p>OVERSIZE 128M,IPCSIZE 2M,ALIGN 也是 2M。</p>
<p><code>oom_check()</code> 就是查询了上面提到的记录着进程信息的共享内存区域,获取当前设备的显存用量,加上请求分配的显存值,和设定的限制值做比较。如果超过限制,就尝试清理下已经结束的进程的显存记录,然后重新计算一遍。</p>
<p>剩下的接口就都是对 cuda 显存分配的一些抽象,把分配结果记录到上面创建的全局列表里面。他们底层又对应着 <code>add_chunk_async()</code></p>
<pre><code class="language-cpp">int allocate_raw(CUdeviceptr *dptr, size_t size);
int free_raw(CUdeviceptr dptr);
int add_chunk_only(CUdeviceptr address,size_t size);
int allocate_async_raw(CUdeviceptr *dptr, size_t size, CUstream hStream); // 基于 add_chunk_async
int free_raw_async(CUdeviceptr dptr, CUstream hStream);
</code></pre>
<p><code>check_memory_type()</code> 就是在 <code>device_overallocated</code> 里面检查有没有查询的指针,判断是设备地址还是 host 侧地址。值得注意的是按这个实现,<code>cuMemAllocManaged</code> 分配出来的地址算到了设备地址里面。</p>
<h2>src/cuda/</h2>
<p><a href="http://libcuda.so">libcuda.so</a> 库劫持逻辑,好多 API 其实没实现劫持方案,只是打印了一下日志,感觉是之后打算做。劫持的时候注意一下 <code>cuGetProcAddress</code> 即可。</p>
<h3>memory.c</h3>
<p>劫持了显存分配 API,逻辑上就是先调用 <code>oom_check()</code> 检查一下,如果超过显存限制,就不分配了直接返回 OOM。</p>
<p>比较特殊的是 <code>cuMemHostAlloc</code>、<code>cuMemAllocHost_v2</code>、<code>cuMemHostRegister_v2</code>,这三个 API,则是先真实触发进行分配,再调用 <code>oom_check()</code> 检查显存,如果超了,就释放掉回滚刚刚到真实分配。这个逻辑感觉很迷惑,这三个 API 的分配并没有涉及到对 <code>src/allocator/allocator.c</code> 中的全局链表的修改,也就是说分配前后实际上检查的效果是一样的,为什么要先分配再回滚呢?</p>
<h2>src/utils.c</h2>
<p>定义了一个跨进程的锁,<code>"/tmp/vgpulock/lock”</code>,<code>try_lock_unified_lock</code> 通过标记 <code>O_EXCL</code> 互斥地打开该文件作为锁。</p>
<p><code>parse_cuda_visible_env</code> 查看环境变量中 <code>CUDA_VISIBLE_DEVICES</code>,尝试对卡的序号进行修正,存到 <code>cuda_to_nvml_map</code> 里面。那这里实际上只支持 nvidia-container-toolkit 对应的 runc 场景,其他厂商都还需要适配。比如昇腾,他们接入 HAMI,是也拿 <code>CUDA_VISIBLE_DEVICES</code> 去做自己的 runc 配置的环境变量了么。</p>
<p><code>mergepid</code> 接收两个 <code>nvmlDeviceGetComputeRunningProcesses</code> 采集到的进程组,合并到一个里面。实际上不要这个函数也行,反正现在只支持单卡,这个函数只是为了对多卡的<code>nvmlDeviceGetComputeRunningProcesses</code> 返回结果做聚合。</p>
<p><code>getextrapid</code> 比较两个<code>nvmlDeviceGetComputeRunningProcesses</code> 采集到的进程组,找到新增的那一个进程。</p>
<p><code>set_task_pid</code> 是为容器内 pid 和 host 侧 pid 建立关联,因为 <code>nvmlDeviceGetComputeRunningProcesses</code> 可以获取到占用 GPU 设备的 host 侧进程号。先调用一次<code>mergepid</code> 将所有占用 GPU 的进程记录到 <code>pre_pids_on_device</code>。<code>cuDevicePrimaryCtxRetain</code> 之后(看上去没激活 ctx 的话不会被 nvml 检测到?),再调用一次 <code>mergepid</code>把所有占用 GPU 的进程记录到 <code>pids_on_device</code>,然后通过<code>getextrapid</code>过滤出来一个新增的进程,就是当前进程在 host 侧的进程号,然后它通过 <code>set_host_pid</code> 把这个 hostpid 写入到 region_info 的共享内存当前进程的分区中。如果有恰好在查找过程中退出的进程,没有影响,因为我们只关心新增的进程,而别的新增进程还卡在 <code>try_lock_unified_lock</code> 那里。</p>
<h2>include</h2>
<h3>libnvml_hook.h</h3>
<p>定义了宏如 NVML_OVERRIDE_CALL,和用于标识 NVML API 的枚举 NVML_OVERRIDE_ENUM_t。实现上不够简洁很多地方可以 #include 同一个 API 列表去做替换。也没看到生成这些头文件的相关脚本,后面升级更新 API 列表很麻烦。</p>
<h3>libcuda_hook.h</h3>
<p>类似 libnvml_hook.h</p>
<h1>HAMi</h1>
<p><a href="https://github.com/Project-HAMi/HAMi">https://github.com/Project-HAMi/HAMi</a></p>
<p>update: 2025-01-05: cbccbd469c636870655dff5f6d2707e206d00a02,更新了 nvidia MIG 和 Metax GPU,加了一些 UT 和文档,特别是有个 <a href="https://github.com/Project-HAMi/HAMi/blob/master/docs/mind-map/HAMI-VGPU-mind-map-Chinese.png">脑图</a> 画得很好。</p>
<p>下面的代码基于 66cabbfac0aebd4ccf19a2d0850c1a2d682b3159</p>
<p>HAMI-core 中的劫持库,会被编译成 <a href="http://libvgpu.so">libvgpu.so</a>,通过挂载 ld.so.preload 文件的方式注入到容器里面做 cuda/nvml 劫持。</p>
<h2>cmd/vGPUmonitor/</h2>
<p>看代码像是内部删了一些东西才开源的,好多没用到的符号……main.go 中开了两个协程分别运行 <code>initMetrics</code> 和 <code>watchAndFeedback</code>。</p>
<h3>metrics.go</h3>
<p>将自定义的 vGPU 数据格式转换收集到 Promethus。</p>
<p><code>initMetrics</code> 监听了 <code>9394</code> 端口,<code>/metrics</code> 会被路由到 Prometheus 服务。由 <code>ClusterManagerCollector</code> 实现 Prometheus 的 <code>Collector</code> 接口,负责采集 metrics 并注册到 Prometheus,注册时定义了 metric 的标签 zone为 <code>vGPU</code>。接口 <code>Collect()</code> 实现时候就是通过 nvml 获取了 host 侧 GPU 指标,通过 <code>ContainerLister.ListContainers()</code> 获取了每个容器的 vGPU 指标。</p>
<h4>testcollector/main.go</h4>
<p>验证 metrics.go 中的 Prometheus 数据采集。</p>
<h3>validation.go</h3>
<p><code>ValidateEnvVars</code> 检查了一下 <code>HOOK_PATH</code> 环境变量是否配置了。</p>
<h3>feedback.go</h3>
<p><code>watchAndFeedback</code> 中每五秒,通过 <code>pkg/monitor/nvidia.ContainerLister</code> 遍历一遍所有容器,记录容器配置的优先级,综合各个容器内的 <code>Priority</code>、<code>RecentKernel</code>、<code>UtilizationSwitch</code> 信息分别修改他们的配置。<br>
这里没看懂修改这三个变量的逻辑,还得回到 HAMI-core 那边联合起来看下。看上去 <code>Priority</code> 是数值越小优先级越高。</p>
<h3>noderpc/noderpc.proto</h3>
<p>定义了一个 gRPC 服务用来获取各个 POD 中的 vGPU 使用情况,<code>rpc GetNodeVGPU (GetNodeVGPURequest) returns (GetNodeVGPUReply) {}</code>。响应中的 <code>sharedRegionT</code>,也就是 HAMI-core 中存放资源切分数据的共享内存中的数据。</p>
<h2>pkg/monitor/</h2>
<h3>nvidia/cudevshr.go</h3>
<p>为 HAMI-core 中共享内存区域内的数据,定义了 v0 和 v1 两个版本的统计信息,通过统一接口 <code>UsageInfo</code>。单个容器的统计信息如下</p>
<pre><code class="language-go">type ContainerUsage struct {
PodUID string
ContainerName string
data []byte
Info UsageInfo
}
</code></pre>
<p><code>ContainerUsage</code> 数据存储在 <code>ContainerLister</code> 类型中,<code>ContainerLister.ListContainers()</code>在上面看到过的 <code>vGPUmonitor</code> 中被用来获取容器内的统计信息。<br>
<code>ContainerLister.Update()</code> 则是遍历各个,通过 <code>loadCache()</code> 函数获取容器的统计数据 <code>ContainerUsage</code>。如果容器内没有调用过 cuInit,那 <code>loadCache()</code> 不会统计到它。<br>
函数<code>loadCache()</code> 实现的逻辑是,查询文件 <code>$HOOK_PATH/containers/$POD_NAME/.cache</code>(<a href="http://libvgpu.so">libvgpu.so</a> 也在该目录内),然后直接 mmap 读取出来转换成符合 <code>UsageInfo</code> 接口的数据。<br>
`</p>
<h2>cmd/scheduler/</h2>
<p>节点/GPU 调度器,实现在 <code>pkg/scheduler/scheduler.go</code>。</p>
<h3>main.go</h3>
<p>启动两个协程运行收集集群中的节点设备信息的 <code>Scheduler.RegisterFromNodeAnnotations()</code> 和采集上报 Prometheus metric 的 <code>initMetrics()</code>,然后默认监听 <code>8080</code> 端口,为 <code>pkg/scheduler/routes/route.go</code> 中定义的 http 服务提供扩展 k8s 调度器服务。</p>
<h3>metrics.go</h3>
<p>和 vGPUmonitor 中的 <code>cmd/vGPUmonitor/metrics.go</code> 一致,不同的地方是容器内信息来自 <code>pkg/scheduler/scheduler.go</code> 中的 <code>Scheduler</code> 类型,例如这里获取 host 侧指标是通过 <code>Scheduler.InspectAllNodesUsage()</code>,每个容器的信息是通过 <code>podManager.GetScheduledPods()</code></p>
<h2>pkg/scheduler</h2>
<p>通过 <code>k8s.io/kube-scheduler/extender/v1</code> API 拓展 k8s 的调度器。k8s 调度器负责将 Pod 分配到合适的节点,而扩展调度器可以让用户自定义调度逻辑,一般都会包含过滤、打分、绑定等机制。</p>
<h3>routes/route.go</h3>
<p><code>cmd/scheduler/main.go</code> 中定义的路由为</p>
<ul>
<li><code>/filter</code> 对应 <code>PredicateRoute</code>。调用 <code>Scheduler.Filter()</code>,处理 Pod 调度过滤逻辑。</li>
<li><code>/bind</code> 对应 <code>Bind</code>,调用<code>Scheduler.Bind()</code>,处理 Pod 调度绑定逻辑。</li>
<li><code>/webhook</code> 对应 <code>WebHookRoute</code>,调用 <code>Scheduler.NewWebHook()</code> 创建 webhook,在 webhook 上调用 <code>ServeHTTP()</code>。</li>
<li><code>/healthz</code> 对应 <code>HealthzRoute</code>,心跳包。</li>
</ul>
<h3>scheduler.go</h3>
<p>调度器类定义</p>
<pre><code class="language-go">type Scheduler struct {
nodeManager
podManager
stopCh chan struct{}
kubeClient kubernetes.Interface
podLister listerscorev1.PodLister
nodeLister listerscorev1.NodeLister
//Node status returned by filter
cachedstatus map[string]*NodeUsage
nodeNotify chan struct{}
//Node Overview
overviewstatus map[string]*NodeUsage
eventRecorder record.EventRecorder
}
</code></pre>
<p>调度器类实现了接口 <code>onUpdateNode()</code>、<code>onDelNode()</code>、<code>onAddNode()</code>,当集群中的节点变更时,会通过回调写入事件到 <code>Scheduler.nodeNotify</code>。</p>
<p>它也实现了<code>onAddPod()</code>、<code>onUpdatePod()</code>、<code>onDelPod()</code>,其中 <code>onUpdatePod()</code> 算是 <code>onAddPod()</code> 的一种特殊情况。这几个回调会直接调用 <code>Scheduler.addPod()</code> 和 <code>Scheduler.delPod()</code>。</p>
<p><code>Scheduler.RegisterFromNodeAnnotations()</code> 会启动一个无限循环,每隔一段时间或收到节点变更通知时,执行节点注册逻辑,直到<code>Scheduler.stopCh</code> 收到信号终止。它通过 <code>Scheduler.nodeLister</code> 获取所有节点,做一次健康检测,然后把每个正常的节点信息建模成 <code>util.NodeInfo</code> 类型,通过 <code>Scheduler.addNode()</code> 函数在调度器中存储节点。注册完毕后会调用一次 <code>Scheduler.getNodesUsage()</code> 获取所有节点/Pod 的设备和显存信息,更新到 <code>Scheduler.cachedstatus</code> 中,不过这个成员看上去还没有怎么用到,可能是之后想做缓存。</p>
<p><code>Scheduler.Filter()</code> 和 <code>Scheduler.Bind()</code> 就是之前 <code>routes/route.go</code> 部分提过的调度器做设备分配和过滤的 http 接口的实现。<code>Scheduler.Filter()</code> 首先尝试清理掉当前请求相关的 Pod 以避免对节点资源统计数据造成偏差,然后通过<code>Scheduler.getNodesUsage()</code>获取当前所有的节点及其资源信息,然后通过 <code>Scheduler.calcScore()</code> 计算一遍节点的分值和对应的节点信息、节点设备信息,按分值排序后获取最优的一个可分配节点,最后调整 Pod 的 annotation,调用 <code>Scheduler.addPod()</code>记录。<code>Scheduler.Bind()</code> 的实现比较简单,就是调用 k8s 的API 获取 node 和 pod,对 node 加锁(锁实现见 <code>pkg/util/nodelock/nodelock.go</code>),然后调用 k8s 的 Bind API 去把 pod 调度到请求的节点上。</p>
<h3>events.go</h3>
<p>基于 <code>k8s.io/client-go/tools/record</code> 中的 <code>EventRecorder</code>,将Pod调度过程中的绑定/过滤事件记录到 k8s事件系统中,便于后续的故障排查和状态监控。</p>
<h3>nodes.go</h3>
<p><code>nodeManager</code> 持有一个节点和设备的哈希表,可以通过 <code>nodeManager.addNode()</code> 和 <code>nodeManager.rmNodeDevice()</code> 添加、删除调度器自身维护的节点上的 GPU 设备信息。</p>
<h3>pods.go</h3>
<p><code>podManager</code> 持有一个已被调度的 Pod 的哈希表,可以通过 <code>podManager.addPod()</code> 和 <code>podManager.delPod()</code> 添加、删除调度器自身维护的 Pod 信息。</p>
<h3>webhook.go</h3>
<p>k8s 允许集群中存在多个调度器。默认情况下,Pod 使用的是 kube-scheduler 调度器。通过设置 SchedulerName 字段,我们可以指定哪个调度器来调度特定的 Pod。<br>
这个 webhook 的 <code>Handle()</code> 就是用来为合法的 Pod 选择使用 HAMI 实现的调度器进行调度。它的实现逻辑是,先检查是否 Pod 内有容器,如果有,再去看是不是特权容器,如果不是特权容器,再调用 <code>pkg/device/devices.go</code> 中定义的 <code>Device</code> 公共接口,检查容器的资源限制、annotation 等是否符合对应设备的配置规范,如果合规就修改调度器。</p>
<h3>score.go</h3>
<p>实现 k8s 中调度算法的核心,打分机制。该文件内的函数之间的调用链为 <code>Scheduler.calcScore()->fitInDevices()->fitInCertainDevice()</code>,对每个节点、每个请求进行遍历,又会对请求中的每种设备需求进行检查,最后返回一个包含节点、设备(包含了针对请求的分配情况)、分值的 <code>policy.NodeScore</code> 的列表。最外层的 <code>Scheduler.calcScore()</code> 是被 <code>Scheduler.Filter()</code> 用在了调度器的节点过滤逻辑中,用来计算节点、设备的分数选择节点设备。</p>
<h3>node_policy.go</h3>
<p>节点调度策略,目前实现了两个,<code>binpack</code> 和 <code>spread</code>,默认为 <code>binpack</code>,优先占满节点,可以通过 POD 级别的 annotation <code>hami.io/node-scheduler-policy</code>进行修改。</p>
<p>为 <code>NodeScoreList</code> 定义了 <code>Less</code> 接口,按节点的分数进行排序。根据节点上的设备使用占比、设备核心使用占比,显存使用占比,三者求和计算分数。</p>
<h3>policy/gpu_policy.go</h3>
<p>GPU 调度策略,目前实现了两个,<code>binpack</code> 和 <code>spread</code>,默认为 <code>spread</code>,优先均匀分布,可以通过 POD 级别的 annotation <code>hami.io/gpu-scheduler-policy</code> 进行修改。</p>
<p>为 <code>DeviceUsageList</code> 定义了 <code>Less</code> 接口,按节点的分数进行排序。根据设备上的设备使用占比、设备核心使用占比,显存使用占比,三者求和计算分数,这里注意是要加上申请的额度的。</p>
<h2>cmd/device-plugin/nvidia</h2>
<p>为英伟达设备实现的 <a href="https://kubernetes.io/docs/concepts/extend-kubernetes/compute-storage-net/device-plugins/">k8s Device Plugin</a>,HAMI 一开始只支持 nvidia,放在这里也算是一个实现范例,其他厂商的 DP 在别的仓库。<br>
也相当于是 nvidia 官方的 <a href="https://github.com/NVIDIA/k8s-device-plugin">DP</a> 的一个拓展。</p>
<h3>main.go</h3>
<p>通过 <code>startPlugins()</code> 启动 DP 的服务,它的主要逻辑如下</p>
<pre><code class="language-go">func startPlugins(c *cli.Context, flags []cli.Flag, restarting bool) ([]plugin.Interface, bool, error) {
config, err := loadConfig(c, flags)
disableResourceRenamingInConfig(config)
devConfig, err := generateDeviceConfigFromNvidia(config, c, flags)
// Update the configuration file with default resources.
err = rm.AddDefaultResourcesToConfig(&devConfig)
// Get the set of plugins.
pluginManager, err := NewPluginManager(&devConfig)
plugins, err := pluginManager.GetPlugins()
// Loop through all plugins, starting them if they have any devices to serve.
for _, p := range plugins {
if len(p.Devices()) == 0 {
continue
}
if err := p.Start(); err != nil {
return plugins, true, nil
}
}
return plugins, false, nil
}
</code></pre>
<p><code>disableResourceRenamingInConfig()</code> 中禁用了官方 DP 中对设备的重命名,之后会恢复回来,应该是受限于当前的 HAMI-core 的实现?</p>
<h3>plugin-manager.go</h3>
<p><code>NewPluginManager()</code> 根据配置生成一个 DP 的工厂用来构建各种依赖的 DP,后面就会看到不只有一个,例如 nvml 也需要单独的 DP 配置。</p>
<h3>vgpucfg.go</h3>
<p>实现了用来解析 HAMI 自定义的参数的工具函数 <code>generateDeviceConfigFromNvidia()</code>。</p>
<h2>pkg/device-plugin/nvidiadevice/nvinternal</h2>
<p>也是基于英伟达官方的 <a href="https://github.com/NVIDIA/k8s-device-plugin/blob/main/internal">https://github.com/NVIDIA/k8s-device-plugin/blob/main/internal</a> 中的 DP 实现进行了拓展。</p>
<h3>plugin</h3>
<h4>api.go</h4>
<p>定义接口</p>
<pre><code class="language-go">type Interface interface {
Devices() rm.Devices
Start() error
Stop() error
}
</code></pre>
<h4>server.go</h4>
<p>定义了类型 <code>NvidiaDevicePlugin</code>,实现下面的 DP 的标准服务接口。</p>
<pre><code>service DevicePlugin {
// GetDevicePluginOptions returns options to be communicated with Device Manager.
rpc GetDevicePluginOptions(Empty) returns (DevicePluginOptions) {}
// ListAndWatch returns a stream of List of Devices
// Whenever a Device state change or a Device disappears, ListAndWatch
// returns the new list
rpc ListAndWatch(Empty) returns (stream ListAndWatchResponse) {}
// Allocate is called during container creation so that the Device
// Plugin can run device specific operations and instruct Kubelet
// of the steps to make the Device available in the container
rpc Allocate(AllocateRequest) returns (AllocateResponse) {}
// GetPreferredAllocation returns a preferred set of devices to allocate
// from a list of available ones. The resulting preferred allocation is not
// guaranteed to be the allocation ultimately performed by the
// devicemanager. It is only designed to help the devicemanager make a more
// informed allocation decision when possible.
rpc GetPreferredAllocation(PreferredAllocationRequest) returns (PreferredAllocationResponse) {}
// PreStartContainer is called, if indicated by Device Plugin during registration phase,
// before each container start. Device plugin can run device specific operations
// such as resetting the device before making devices available to the container.
rpc PreStartContainer(PreStartContainerRequest) returns (PreStartContainerResponse) {}
}
</code></pre>
<p><code>cmd/device-plugin/nvidia/main.go</code> 通过它的 <code>NvidiaDevicePlugin.Start()</code> 函数拉起 DP 服务 <code>NvidiaDevicePlugin.Serve()</code>,通过 <code>NvidiaDevicePlugin.Register()</code> 注册自身到 kubelet,然后开启一个线程调用 <code>ResourceManager.CheckHealth()</code> 持续做健康检查,一个线程调用 <code>NvidiaDevicePlugin.WatchAndRegister()</code> 定期更新节点设备信息。</p>
<p><code>NvidiaDevicePlugin.Serve()</code> 就是启动了一个 uds server,提供上面的 <code>DevicePlugin</code> service。</p>
<p><code>DevicePlugin.PreStartContainer()</code> 和 <code>DevicePlugin.GetPreferredAllocation()</code> 都是空实现,<code>DevicePlugin.ListAndWatch()</code> 在设备有健康状态变化的时候,返回 <code>ResourceManager.Devices().GetPluginDevices()</code> 的结果。</p>
<p><code>DevicePlugin.Allocate()</code> 先调用 <code>util.GetNextDeviceRequest()</code> 从 pending 状态的 Pod 中找到需要 GPU 设备的容器;调用 <code>NvidiaDevicePlugin.getAllocateResponse()</code> ,其中根据 <code>NvidiaDevicePlugin.deviceIDsFromAnnotatedDeviceIDs()</code> 根据请求信息中的设备获取符合 CDI 规范的设备 ID,然后在 <code>NvidiaDevicePlugin.getAllocateResponseForCDI()</code> 中进一步根据 CDI 规范,修改响应的格式以触发 CDI;然后修改容器的环境变量和 mount,在这里配置 vgpu/HAMi-core(这里一些操作不可以放到 <code>DevicePlugin.PreStartContainer()</code> 里面吗)。</p>
<h4>register.go</h4>
<p>DP 的 grpc 服务器启动后,会单独启动一个线程调用该文件内的 <code>WatchAndRegister()</code>,定期获取节点侧设备信息,更新到节点的 annotation 中。</p>
<h4>manager</h4>
<p>不同平台下的 DP manager。</p>
<h5>api.go</h5>
<p>定义接口</p>
<pre><code class="language-go">type Interface interface {
GetPlugins() ([]plugin.Interface, error)
CreateCDISpecFile() error
}
</code></pre>
<h5>factory.go</h5>
<p><code>manager</code> 类型用来初始化 nvml、cdi,解析环境是 nvml 类型还是 tegra 类型,与后面的 manager 实现、<code>rm</code> 模块有关。</p>
<p><code>New()</code> 中 <code>manager</code> 被拓展成了 <code>nvmlmanager</code>(一般情况都是基于 nvml 来管理)、<code>tegramanager</code> (tegra 设备使用)和 <code>null</code>(错误情况下的 fallback) 三类 manager,他们均需要实现 <code>api.Interface</code>。</p>
<h5>null.go</h5>
<p><code>null</code> manager 实现。</p>
<h5>nvml.go</h5>
<p><code>nvmlmanager</code> 实现。<code>nvmlmanager.GetPlugins()</code> 接口,通过 <code>rm.NewNVMLResourceManagers()</code> 获取所有资源,对每个资源,通过 <code>plugin/server.go</code> 中 <code>plugin.NewNvidiaDevicePlugin()</code> 构建 DP。</p>
<h5>tegra.go</h5>
<p><code>tegramanager</code> 实现。</p>
<h3>rm</h3>
<p>分配、管理、监控每个资源对应的 GPU 设备。</p>
<h4>rm.go</h4>
<p>实现 <code>resourceManager</code>,负责管理 GPU 设备。定义接口 <code>ResourceManager</code>。</p>
<pre><code class="language-go">type ResourceManager interface {
Resource() spec.ResourceName
Devices() Devices
GetDevicePaths([]string) []string
GetPreferredAllocation(available, required []string, size int) ([]string, error)
CheckHealth(stop <-chan interface{}, unhealthy chan<- *Device) error
}
</code></pre>
<p><code>NewResourceManagers()</code> 为每个资源创建 <code>ResourceManager</code> 接口类型,一般来说使用的是 <code>NewNVMLResourceManagers()</code> (不需要考虑 tegra 设备)。</p>
<h4>allocate.go</h4>
<p>实现了两个 GPU 分配算法,用户可以使用 <code>resourceManager.getPreferredAllocation()</code> 获取分配出的 GPU 设备。</p>
<p>其中一个集成了 <a href="https://github.com/NVIDIA/go-gpuallocator/">https://github.com/NVIDIA/go-gpuallocator/</a> 中的 GPU 分配器,它会借助 nvml 识别拓扑关系,按预定的策略选择合适的 GPU 设备,<code>resourceManager.alignedAlloc()</code>。<br>
另一个则是考虑了过往的分配情况,尽可能均匀地完成分配,<code>resourceManager.distributedAlloc()</code>。</p>
<h4>devices.go</h4>
<p><code>Device</code> 类包装 k8s 的 DP 中对设备的抽象,即 <a href="http://k8s.io/kubelet/pkg/apis/deviceplugin/v1beta1">k8s.io/kubelet/pkg/apis/deviceplugin/v1beta1</a> 中的 <code>Device</code>。提供了一组公共的接口</p>
<pre><code class="language-go">type deviceInfo interface {
GetUUID() (string, error)
GetPaths() ([]string, error)
GetNumaNode() (bool, int, error)
}
</code></pre>
<p><code>Devices.GetPluginDevices()</code> 会被 DP 的 <code>ListAndWatch()</code> 接口使用。</p>
<h4>device_map.go</h4>
<p><code>DeviceMap</code> 基于给定的 libnvml、资源名、nvidia 官方 DP 的配置,构建资源名到 HAMI <code>resourceManager</code> 中的设备抽象的映射(device.go 中的 <code>Device</code> 类型)。</p>
<h4>health.go</h4>
<p>检查 GPU 设备的监控状态,允许通过环境变量 <code>DP_DISABLE_HEALTHCHECKS</code> 指定一些可忽略的 xid 错误。目前默认忽略下面的 xid,因为他们只表明用户应用出错了但是设备可能仍然可用。</p>
<pre><code class="language-go">// http://docs.nvidia.com/deploy/xid-errors/index.html#topic_4
// Application errors: the GPU should still be healthy
applicationErrorXids := []uint64{
13, // Graphics Engine Exception
31, // GPU memory page fault
43, // GPU stopped processing
45, // Preemptive cleanup, due to previous errors
68, // Video processor exception
}
</code></pre>
<h4>nvml_devices.go</h4>
<p><code>nvmlDevice</code> 和 <code>nvmlMigDevice</code> 类型包装了 <a href="http://github.com/NVIDIA/go-nvlib/pkg/nvml">github.com/NVIDIA/go-nvlib/pkg/nvml</a> 的 <code>Device</code>。同样是实现了 <code>deviceInfo</code> 接口。</p>
<h4>nvml_manager.go</h4>
<p><code>nvmlResourceManager</code> 包装了 <code>resourceManager</code>,集成了 <a href="http://github.com/NVIDIA/go-nvlib/pkg/nvml">github.com/NVIDIA/go-nvlib/pkg/nvml</a> 中的 nvml接口。</p>
<h4>wsl.go</h4>
<p><code>wslDevice</code> 包装了一层 <code>nvmlDevice</code>。= = 还支持 wsl 的……</p>
<h4>tegra_devices.go、tegra_manager.go</h4>
<p>Tegra 设备只支持 <code>resourceManager.distributedAlloc()</code> 分配策略。</p>
<p><code>tegraResourceManager</code> 包装了 <code>resourceManager</code>。</p>
<p>= = 还支持 tegra 的……</p>
<h3>cdi</h3>
<p>借助官方实现 <code>github.com/NVIDIA/nvidia-container-toolkit/pkg/nvcdi</code>,为 DP 使用的 nvidia 设备创建 CDI specs</p>
<h4>api.go</h4>
<p>定义实现 CDI 的接口</p>
<pre><code class="language-go">type Interface interface {
CreateSpecFile() error
QualifiedName(string, string) string
}
</code></pre>
<h4>factory.go</h4>
<p>CDI <code>Interface</code> 的工厂函数,如果没有检测到 nvidia 设备,就创建一个空实现,也就无法生成 CDI specs。</p>
<h4>cdi.go</h4>
<p>定义一个 <code>cdiHandler</code> 类型用于实现生成 CDI 的接口 <code>Interface</code></p>
<pre><code class="language-go">type cdiHandler struct {
logger *logrus.Logger
nvml nvml.Interface // github.com/NVIDIA/go-nvlib/pkg/nvml
nvdevice nvdevice.Interface // github.com/NVIDIA/go-nvlib/pkg/nvlib/device
driverRoot string
targetDriverRoot string
nvidiaCTKPath string
cdiRoot string
vendor string
deviceIDStrategy string
enabled bool
gdsEnabled bool // GPUDirect Storage
mofedEnabled bool // Mellanox OpenFabrics Enterprise Distribution
cdilibs map[string]nvcdi.Interface // github.com/NVIDIA/nvidia-container-toolkit/pkg/nvcdi
}
</code></pre>
<h3>mig/mig.go</h3>
<p><code>GetMigCapabilityDevicePaths</code> 获取 nvidia MIG 切分模式下的设备文件。</p>
]]></content:encoded>
<author>peihao.young@gmail.com (Peihao Yang)</author>
</item>
<item>
<title><![CDATA[torch 源码阅读]]></title>
<link>https://forsworns.github.io///zh/blogs/20241008/</link>
<guid>https://forsworns.github.io///zh/blogs/20241008/</guid>
<pubDate>Tue, 08 Oct 2024 00:00:00 GMT</pubDate>
<description><![CDATA[Pytorch2 recap]]></description>
<content:encoded><![CDATA[<p>[[toc]]</p>
<h2>pytorch2 中打开日志</h2>
<p>import torch<br>
import logging<br>
torch._logging.set_logs(all=logging.DEBUG)</p>
<h2>nvrtc</h2>
<p>torch.compile 时大量的子进程占用 GPU 设备。已知单纯调用 libnvrtc 和 libnvJitLink 不会引用 GPU 设备。</p>
<pre><code>import torch
import logging
torch._logging.set_logs(all=logging.DEBUG)
# “reduce-overhead” 模式才会用到 cuda graph
@torch.compile(mode="reduce-overhead")
def my_model(x):
y = torch.matmul(x, x).cuda()
# side effect breaks graph construction
input("during capture")
y = torch.matmul(y, x).cuda()
return y
x = torch.randn(10, 10).cuda()
print("graph exec 1", flush=True)
y = my_model(x)
print("graph exec 2", flush=True)
y = my_model(x)
print("y", y, flush=True)
</code></pre>
<p>Pytorch 中通过宏 AT_CUDA_NVRTC_CHECK 调用 nvrtc 的代码处理错误。<a href="http://libnvrtc.so">libnvrtc.so</a> 库中的符号类似 <a href="http://libcuda.so">libcuda.so</a> 等是通过下面的函数动态获取的</p>
<pre><code class="language-c++">const at::cuda::NVRTC& nvrtc() {
return at::globalContext().getNVRTC();
}
</code></pre>
<p>CUDAHook 和 nvrtc 的具体的实现是在<br>
<code>aten/src/ATen/cuda/detail/CUDAHooks.cpp</code><br>
和<br>
<code>aten/src/ATen/cuda/detail/LazyNVRTC.cpp</code>。在 <code>aten/src/ATen/cuda/detail/LazyNVRTC.cpp</code> 中,<code>lazyNVRTC</code> 这个全局变量上除了 nvrtc 的 API,还会通过 <code>dlsym</code> 懒加载 <a href="http://libcuda.so">libcuda.so</a> 中的部分 API。 可能就是这里 dlopen 的 <a href="http://libcuda.so">libcuda.so</a> 仍占用设备。但是这里也说不通,cuInit 和 CUcontext 相关的函数不可能不先调用。</p>
<p>除了下面 torch.compile 相关的会调用到 nvrtc 编译 cuda 算子,aten 库中例如 <code>aten/src/ATen/native/cuda/jit_utils.cpp</code> 的 <code>jit_pwise_function</code> 也调用了 nvrtc,但是是用来做一些循环展开等算子优化的。</p>
<h2>torch.compile 过程中的 nvrtc 调用</h2>
<p>被用于代码生成</p>
<p>torch/_dynamo/trace_rules.py 中定义了一系列指导前端的 trace rule,其中的 “torch._C._te.construct_codegen” 是定义在 torch/csrc/jit/tensorexpr/tensorexpr_init.cpp 中的 <code>construct_codegen -> CudaCodeGen</code>,<code>CudaCodeGen::Initialize -> CudaCodeGen::CompileToNVRTC</code> 可以调用到 nvrtc 进行即时编译。</p>
<p>将 CudaAnalysis 分析出的代树,通过 CudaPrinter、GPUMetaVarRewriter 辅助结构体交给其他辅助函数做重写,结果写入到一个 ostringstream,交给 nvrtc 编译。</p>
<p>被用于算子融合</p>
<p>prim::FusionGroup (torch/csrc/jit/runtime/register_prim_ops_fulljit.cpp) -> runFusion -> launchFusion(torch/csrc/jit/codegen/fuser/executor.cpp) -> launch_raw<br>
compileKernel (torch/csrc/jit/codegen/fuser/compiler.cpp) 调用了 getConstructor 拿到 registerFusionBackend 中注册到全局的 <code>createFusionKernel</code>(torch/csrc/jit/codegen/fuser/cuda/fused_kernel.cpp 或 <code>torch/csrc/jit/codegen/fuser/cpu/fused_kernel.cpp</code>) 。</p>
<h2>torch.compile 其他 API</h2>
<h3>前端 dynamo</h3>
<p>torch/_dynamo/eval_frame.py 中的 <code>_optimize()</code> 前端 TorchDynamo 的入口函数。</p>
<p>torch/_dynamo/convert_frame.py 定义了 <code>ConvertFrame</code>、 <code>ConvertFrameAssert</code> 等仿函数类,最终调用同文件内的 <code>_compile</code> 函数,将栈帧转换成 FX graph。</p>
<h3>后端 inductor</h3>
<p>torch/_inductor/compile_fx.py 中的 <code>compile_fx</code>,在 torch/_dynamo/backends/inductor.py 里面通过 <code>register_backend</code> 注册到全局的后端 fx graph 入口函数。</p>
<p>torch/_inductor/async_compile.py 异步编译。维护了一个进程池,例如 <code>triton</code> 后端就通过下面这个函数,向进程池提交了一个编译任务然后返回了一个 future。</p>
<pre><code class="language-python">def triton(self, kernel_name: str, source_code: str, device_str: str = "cuda"):
kernel = TritonCodeCache.load(kernel_name, source_code)
return TritonFuture(
kernel,
self.process_pool().submit(
_worker_compile_triton,
kernel._reload_in_subproc,
extra_env,
),
)
</code></pre>
<p>所以能看出,多出来的那些进程,是通过这里生成的。可以通过下面的 <code>TORCHINDUCTOR_COMPILE_THREADS</code> 环境变量修改。<br>
然后进程池是使用的 concurrent.futures.ProcessPoolExecutor,<strong>基于 fork,所以可能 cuda 没有用到,但是 /dev/nvidia0 等 fd 也被占用了</strong>。</p>
<p>torch/_inductor/codecache.py 代码编译缓存,例如 CUDACodeCache 为 cuda 代码的编译缓存。如果需要,调用 <code>cuda_compile_command</code> 函数进行编译。</p>
<h3>其他</h3>
<p>torch/_inductor/config.py 中,通过 <code>decide_compile_threads</code> 获取了 cpu 核心数。通过环境变量 <code>TORCHINDUCTOR_COMPILE_THREADS</code> 可以修改。</p>
<h2>cuda graph 代码</h2>
<h3>前端 dynamo</h3>
<p>torch/_dynamo/backends/cudagraphs.py 中的 <code>CudagraphsBackend</code>,通过 <code>register_backend</code> 注册到全局的 cuda graph 后端入口。</p>
<h3>在后端 inductor</h3>
<p>torch/_inductor/cudagraph_trees.py<br>
torch/_inductor/cudagraph_utils.py</p>
<h2>torch 日志打印</h2>
<p>pytorch2 中,<br>
环境变量 TORCH_LOGS="+inductor,+dynamo"<br>
或者也可以通过 API 直接设置 torch._logging.set_logs(all=logging.DEBUG)</p>
<h2>c10/cuda/CUDACachingAllocator.cpp</h2>
<p>torch 的两个显存分配器实现,PYTORCH_NO_CUDA_MEMORY_CACHING 控制版本。有一个torch 版本的新分配器依赖 NVML,平台不支持 NVML 的时候可以先用这个环境变量关了新分配器。</p>
<h2>torch/nn/parallel/distributed.py</h2>
<p>完全基于集合通信库的 DP 实现 <code>DistributedDataParallel</code>(DDP),各个 rank 独自持有模型。</p>
<p>首先看 <code>DistributedDataParallel</code> 的构造函数。<code>self._module_parameters</code> 中存了所有未被参数过滤掉的模型参数。</p>
<p>初始化阶段 <code>self._verify_param_shape_across_processes()</code>、<code>self._sync_module_states()</code> 用来在不同 rank 间检查、同步 module 的初始状态,这一步会用 broadcast 集合通信,buffer 大小在代码里写死了是 250 MB。最关键的是 <code>self._ddp_init_helper()</code> 去初始化 reducer,这一步完了基本初始化就结束了,剩下的就是一些混合精度 AMP 相关的代码等,该函数的注释如下,解释得很详细了</p>
<pre><code class="language-python">"""
DDP init helper function to manage parameters, grad hooks, logging, and SyncBatchNorm.
Initialization helper function that does the following:
(1) bucketing the parameters for reductions
(2) resetting the bucketing states
(3) registering the grad hooks
(4) Logging construction-time DDP logging data
(5) passing a handle of DDP to SyncBatchNorm Layer
"""
</code></pre>
<p>DDP 通过将参数划分到不同的桶里面,来实现梯度 reduction 和 反向传播的 overlap 优化,掩盖 reduction 耗时,桶的默认大小是 25MB,可以通过 <code>bucket_cap_mb</code> 参数控制,桶的初始化也在 <code>self._ddp_init_helper()</code> ,调用到的方法是 <code>torch/csrc/distributed/c10d/reducer.cpp</code> 里面定义的 <code>compute_bucket_assignment_by_size</code>,获取每个桶的 idx 和 size。划分好桶之后,再通过 <code>torch/csrc/distributed/c10d/reducer.cpp</code> 中实现的 <code>Reducer</code> 类型的构造函数,完成初始化。</p>
<p><code>compute_bucket_assignment_by_size()</code>的逻辑是:先构造一个桶的哈希表,每个桶内可能有多个张量,哈希表的键是通过张量的数据类型和它所在的设备哈希出来的,张量数据的大小计算方式就是张量规模乘上它的数据类型的大小。当一个键对应的桶被塞满,就要将当前的桶添加到返回列表里面,然后为相应键重建一个桶。最后再把剩余的没满的桶填充到返回列表里面。这个函数的实现上有一个技巧,就是希望尽可能让返回的列表中桶的顺序按模型中张量出现的顺序排列。在没有 torch 上层代码提供提示的情况下,这个函数里面对返回列表中的每个桶里面最小的张量序号进行了排序,假设序号小的张量是优先出现在模型中的参数。回到<code>self._ddp_init_helper()</code>中,它又将 <code>compute_bucket_assignment_by_size()</code> 返回的列表翻转了一下,希望优先处理先被反向传播过程处理到的张量。</p>
<p>前面提到的 <code>torch/csrc/distributed/c10d/reducer.cpp</code> 中的 <code>Reducer</code>,进行 all reduce 通信的方法实际上是 <code>Reducer::run_comm_hook()</code>。它的调用链路是 <code>Reducer::autograd_hook()-> Reducer::mark_variable_ready() -> Reducer::mark_bucket_ready() -> Reducer::all_reduce_bucket() -> Reducer::run_comm_hook()</code>。<code>Reducer::autograd_hook()</code> 被注册给了 <code>torch::autograd::impl::grad_accumulator::add_post_hook()</code>,会在梯度计算完毕累加到了梯度张量后执行,而且只会在 pytorch 的 autograd 线程上执行。</p>
<blockquote>
<p>torch 的 autograd 线程在 torch/csrc/autograd/engine.cpp 的 <code>Engine::start_device_threads()->Engine::thread_init()</code> 创建。</p>
</blockquote>
<pre><code class="language-cpp">c10::intrusive_ptr<c10::ivalue::Future> Reducer::run_comm_hook(
GradBucket& grad_bucket) {
if (comm_hook_ == nullptr) {
// `Reducer` 构造时的参数没有配置过的话,会从 `Reducer::process_group_` 构造一个 `_AllReduceBySumCommHook`,
// 然后对每个 bucket 做 all reduce。
return run_allreduce_hook(grad_bucket);
} else {
return comm_hook_->runHook(grad_bucket);
}
}
</code></pre>
<p>初始化函数里面还有一段有趣的代码是通过 <code>torch._dynamo.config._get_optimize_ddp_mode()</code> 获取了 torch 2 中编译器前端 torchdynamo 的设置,选择启用实验性的 python reducer,而不是用默认的 cpp 实现的 reducer。去读 <code>torch/_dynamo/config.py</code> 的话就会知道,<code>DistributedDataParallel</code> 默认会对 DDP 的通信和计算做 overlap 以掩盖通信开销。但是由于 torch 2 依赖于 PEP 523,所以纯 python 实现的 reducer 更有利于 torch 做自动分析和优化。比如通常路径上 all_reduce 操作调用的是 <code>torch.distributed.distributed_c10d.all_reduce()</code>,但是 python 版的是 <code>torch.distributed._functional_collectives.all_reduce()</code> 中的实现。通过 <code>DistributedDataParallel._get_active_ddp_module()</code> 类方法可以把 DDP 对象暴露给 torchdynamo,<code>DistributedDataParallel._inside_ddp_forward()</code> 这个 contextmanager 则是在调用 <code>DistributedDataParallel._run_ddp_forward()</code> 前就禁用了 torchdynamo。</p>
<p>对于 <code>torch.distributed.distributed_c10d.all_reduce()</code>,它默认调用的是 <code>_get_default_group()</code> 获取的 <code>ProcessGroup</code> 上的 <code>all_reduce()</code> 接口,这实现在 <code>torch/csrc/distributed/c10d/ProcessGroupWrapper.cpp</code> 中的公共抽象 <code>ProcessGroupWrapper</code>,它继承于 <code>Backend</code> 类型。torch 支持了多种集合通信库,每个实现也都继承了 <code>Backend</code>,如 NCCL、GLOO、UCC 等。一般我们只关心 NCCL,它实现在 <code>torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp</code>。</p>
<p>通过 <code>PYTORCH_DDP_USE_SIDE_STREAM</code> 环境变量可以新开一个 cuda steram 做 H2D 的拷贝。</p>
<h2>torch/utils/data/distributed.py</h2>
<p><code>DistributedSampler</code> 配合 <code>DistributedDataParallel</code> 使用,对输入进行分片。<br>
实现很简单,就是在 <code>self.__iter__</code> 函数中 <code>indices[self.rank : self.total_size : self.num_replicas]</code>,对整个 <code>self.total_size</code> 长的数据,间隔 <code>self.num_replicas</code> 按 <code>self.rank</code> 选一个。<code>self.rank</code> 可以通过 <code>dist.get_rank</code> 获取。</p>
<h2>torch/distributed/</h2>
<p><code>DTensor</code> 的各种 TP 方式的实现,和与其他并行方式的集成。DTensor 本身实现在 <code>torch/distributed/_tensor/api.py</code>。</p>
<h3>tensor/parallel/ddp.py</h3>
<p><code>DistributedDataParallel</code> 中调用的 <code>_pre_dp_module_transform()</code> 的实现,便于 DDP 和 TP 结合(torch 的 TP 依赖于 DTensor)。它注册了两个更新 DTensor 的钩子,一个用于在前向传播之前将本地张量转换回 DTensor,另一个用于在前向传播之后将 DTensor 转换回张量。避免 DDP 对 DTensor 参数的特殊处理,并使 DTensor 的梯度能够传递回 DDP 的梯度桶。</p>
<pre><code class="language-python">def _pre_dp_module_transform(module: nn.Module):
_localize_dtensor(module, None, None)
# Recontruct DTensor parameters from local tensors
module.register_forward_pre_hook(_reconstruct_dtensor)
# Convert DTensor parameters to local tensors
module.register_forward_hook(_localize_dtensor)
</code></pre>
<h3>nn</h3>
<p>定义了 torch 用户可以主动使用的集合通信接口 <code>torch.distributed.nn.functional</code>,主动创建位于远端进程的 module <code>torch.distributed.nn.RemoteModule</code>。</p>
<h3>algorithms</h3>
<p>一些分布式下的算法实现。</p>
<p>如 <code>torch/distributed/algorithms/model_averaging/averagers.py</code> 定义了用户可以直接调用的对各个 rank 的参数做均值的 <code>PeriodicModelAverager</code>,可以用于主动同步模型参数、和 PostLocalSGDOptimizer 结合用于优化器等。</p>
<p><code>torch/distributed/optim/</code> 实现分布式的优化器,例如 <code>PostLocalSGDOptimizer</code>、<code>ZeroRedundancyOptimizer</code>。</p>
<p><code>torch/distributed/algorithms/_comm_hooks/default_hooks.py</code> 是分布式训练中默认的 hook,可以通过 <code>model.register_comm_hook</code> 注册别的 hook。如 <code>torch/distributed/algorithms/ddp_comm_hooks/</code> 下的 <code>allreduce_hook()</code>、<code>post_localSGD_hook()</code>。注册 hook 的函数是 <code>DistributedDataParallel.register_comm_hook()</code></p>
<h2>torch/nn/parallel/data_parallel.py</h2>
<p>实现了 <code>DataParallel</code>,仅持有一份模型,每次前向更新都会在不同设备间拷贝需要并行的参数,一般已不使用。</p>
<p><code>torch.nn.parallel.replicate()</code> 用于在各个设备上复制模型,它主要调用了 <code>_broadcast_coalesced_reshape() -> comm._broadcast_coalesced()</code>。<code>broadcast_coalesced()</code> 是个 C 函数,实现在 <code>torch/csrc/cuda/comm.cpp</code>。具体实现在 <code>_broadcast_out_impl</code>,这里通过一个宏控制了是否使用 NCCL,否则就是直接 CUDA D2D 拷贝。</p>
<pre><code class="language-python">static inline std::vector<Tensor>& _broadcast_out_impl(
const Tensor& tensor,
std::vector<Tensor>& out_tensors) {
#ifdef USE_NCCL
std::vector<Tensor> nccl_list;
nccl_list.reserve(out_tensors.size() + 1);
nccl_list.emplace_back(tensor);
for (auto& out_tensor : out_tensors) {
nccl_list.emplace_back(out_tensor);
}
if (nccl::is_available(nccl_list)) {
nccl::broadcast(nccl_list);
} else {
#else
{
#endif
for (auto& out_tensor : out_tensors) {
out_tensor.copy_(tensor, /*non_blocking=*/true);
}
}
return out_tensors;
}
</code></pre>
<p><code>torch.nn.parallel.parallel_apply()</code> 用于在不同设备上并行计算,走得就是在不同线程、不同 <code>torch.cuda.device()、torch.cuda.stream()</code> 下调用 module 的方式。</p>
<h2>torch.randn() 的实现</h2>
<p>randn的具体实现方式<br>
/aten/src/ATen/native/TensorFactories.cpp: Tensor rand() -><br>
/aten/src/ATen/native/Distributions.cpp: Tensor& uniform_() -><br>
/aten/src/ATen/native/DistributionTemplates.h: at::Tensor& uniform_impl_() -><br>
/aten/src/ATen/native/cuda/DistributionUniform.cu: void uniform_kernel() -><br>
/aten/src/ATen/native/cuda/DistributionTemplates.h: void uniform_kernel(), void uniform_and_transform(),void distribution_nullary_kernel()<br>
void uniform_and_transform() 里面根据数据类型,通过 distribution_nullary_kernel()加载了一个遍历 Tensor 的核函数,逐项调用 curand API curand_uniform4 或 curand_uniform2_double进行填充。<br>
动态获取 cuda API 符号<br>
<a href="https://github.com/pytorch/pytorch/blob/main/c10/cuda/driver_api.cpp#L40">https://github.com/pytorch/pytorch/blob/main/c10/cuda/driver_api.cpp#L40</a><br>
在单个单例中,搜索 cuda driver API 和 nvml API</p>
<h2>/aten/src/ATen/native/TensorFactories.cpp</h2>
<p>张量的工厂类</p>
<h2>/aten/src/ATen/native/Convolution.cpp</h2>
<p>卷积实现,<br>
例如正向推理的 cudnn 实现<br>
cudnn_convolution -> cudnn_convolution_forward -> raw_cudnn_convolution_forward_out -> raw_cudnn_convolution_forward_out_32bit -> cudnnConvolutionForward<br>
<a href="https://github1s.com/pytorch/pytorch/blob/v2.1.0/aten/src/ATen/native/cudnn/ConvShared.cpp#L180">https://github1s.com/pytorch/pytorch/blob/v2.1.0/aten/src/ATen/native/cudnn/ConvShared.cpp#L180</a><br>
<a href="https://github1s.com/pytorch/pytorch/blob/v2.1.0/aten/src/ATen/native/cudnn/Conv_v7.cpp#L629">https://github1s.com/pytorch/pytorch/blob/v2.1.0/aten/src/ATen/native/cudnn/Conv_v7.cpp#L629</a><br>
例如反向传播的实现<br>
<a href="https://github1s.com/pytorch/pytorch/blob/v2.1.0/aten/src/ATen/native/Convolution.cpp#L1974-L1978">https://github1s.com/pytorch/pytorch/blob/v2.1.0/aten/src/ATen/native/Convolution.cpp#L1974-L1978</a><br>
这里去根据后端选择对应实现,例如 cudnn_convolution_backward_stub,实现在 /aten/src/ATen/native/cudnn/ConvShared.cpp。</p>
<h2>/aten/src/ATen/native/cuda/CUDALoops.cuh</h2>
<p>借助 cuda 遍历 Tensor 中的每个元素,上述工厂类中会通过 Tensor::fill_()方法调用到它,对张量进行赋值。</p>
<h2>/aten/src/ATen/native/native_functions.yaml</h2>
<p>每个 native 算子有多个后端的 native 实现,该文件描述了这些变体。<br>
例如 fft 变换,/aten/src/ATen/native/SpectralOps.cpp 中的 Tensor stft()调用了对应的 native 算子 _fft_r2c,这又对应了两类后端实现,_fft_r2c_cufft 和 _fft_r2c_mkl</p>
<ul>
<li>func: _fft_r2c(Tensor self, int[] dim, int normalization, bool onesided) -> Tensor<br>
variants: function<br>
dispatch:<br>
CPU: _fft_r2c_mkl<br>
CUDA: _fft_r2c_cufft</li>
</ul>
]]></content:encoded>
<author>peihao.young@gmail.com (Peihao Yang)</author>
</item>
<item>
<title><![CDATA[kTransformers 源码阅读]]></title>
<link>https://forsworns.github.io///zh/blogs/20240924/</link>
<guid>https://forsworns.github.io///zh/blogs/20240924/</guid>
<pubDate>Tue, 24 Sep 2024 00:00:00 GMT</pubDate>
<description><![CDATA[可便捷调优的推理服务器]]></description>
<content:encoded><![CDATA[<p>[[toc]]</p>
<p>基于 <a href="https://github.com/kvcache-ai/ktransformers/tree/0f054fe4ff73133378a8de8ae17f47d5f3ec680a">https://github.com/kvcache-ai/ktransformers/tree/0f054fe4ff73133378a8de8ae17f47d5f3ec680a</a></p>
<p>kvcache-ai/ktransformers 是一个很有趣的推理服务器,只需要写一个 YAML 配置文件,就可以动态地注入借助 SIMD、<a href="https://github.com/IST-DASLab/marlin">marlin</a> 等方法调优后的模块,替换原始模型中的结构。</p>
<p>清华的章明星老师有一个讲座视频:<a href="https://www.bilibili.com/video/BV1CAW6eiENy/">https://www.bilibili.com/video/BV1CAW6eiENy/</a></p>
<p>先从服务端入口读起来</p>
<h2>ktransformers/server/api/<strong>init</strong>.py</h2>
<p>服务端的路由定义,支持多种风格:ollama、OpenAI、web</p>
<h2>ktransformers/server/api/openai/endpoints/chat.py</h2>
<p>OpenAI 接口实现在这里,使用的是 fastapi 构建的服务器,一些用户参数的解析可以在 ktransformers/server/schemas/endpoints/chat.py 里面看到,例如 <code>/chat/completions</code> 的参数 <code>ChatCompletionCreate</code> 现在其实只支持了 steram 类型响应,也忽略了 model 参数。</p>
<h2>ktransformers/server/utils/create_interface.py</h2>
<p>从上面 OpenAI 的 <code>/chat/completions</code> 实现可以继续找到后端推理服务的实现。该文件内实现了一个全局单例 GlobalInterface,server 在处理推理请求的时候都会去获取这个单例,调用 <code>inference</code> 接口执行推理,目前支持两种后端:transformers、ktransformers,都在 ktransformers/server/backend/interfaces 路径下,还有个 <code>exllamav2</code> 看上去还未实现。</p>
<p>ktransformers 和 transformers 后端间是存在继承关系的:<code>KTransformersInterface -> TransformersInterface -> BackendInterfaceBase</code>。他们重要的成员都是 tokenizer、model、cache。只是 <code>KTransformersInterface</code> 里面的 model 是通过 <code>optimize_and_load_gguf()</code> 函数根据 YAML 配置优化过的模型。</p>
<p><code>TransformersInterface</code> 做一次推理的完整的调用链是</p>
<pre><code>inference()
-> prefill()
-> generate()
-> decode_one_tokens()`。
</code></pre>
<p><code>KTransformersInterface</code> 只重载了 <code>TransformersInterface</code> 的 <code>decode_one_tokens()</code> 这个方法,也就是只有 decode 阶段是优化过的,prefill 阶段是默认的。</p>
<h2>ktransformers/models/</h2>
<p>这个目录下面放的都是模型的具体实现。<br>
我们测试的 deepseek 就放在了这里,我们支持了 decode 阶段现的任意保存恢复启停,但是 prefill 阶段做保存恢复就 segmentfault 了。</p>
<h2>ktransformers/operators/</h2>
<p>Ktransformer 的调优实现,即开头提到的注入到模型中的模块。</p>
<h2>ktransformers/util/cuda_graph_runner.py</h2>
<p>一开始关注这个项目,就是听到前面的视频里面,章老师提到用到了他们用了 CUDA Graph,看了下这段用的是 torch.cuda 的接口,实现得很简洁清晰,刚好可以拿来做我们的测试用例。 直接去跑 Llama.cpp 我一直调用不到 CUDA Graph API = =</p>
<p>这里实现的话就是借助了 CUDA API 的自动捕获能力。因为图的输入输出 buffer 都是 capture 得到的,所以需要注意保证在执行 CUDA Graph 的时候还是固定地址,因此可以看到 <code>CUDAGraphRunner::forward</code> 需要把数据拷贝到之前已经创建好的固定的 buffer 里面。</p>
<pre><code class="language-python">class CUDAGraphRunner:
def forward(
self,
cur_token,
position_ids,
cache_position,
) -> torch.Tensor:
# Copy the input tensors to the input buffers.
inputs_embeds = self.model.model.embed_tokens(cur_token.to("cpu"))
self.input_buffers["inputs_embeds"].copy_(inputs_embeds)
self.input_buffers["position_ids"].copy_(position_ids)
self.input_buffers["cache_position"].copy_(cache_position)
# Run the graph.
self.graph.replay()
torch.cuda.synchronize(self.main_device)
# Return the output tensor.
return self.output_buffers["logits"]
</code></pre>
<p><code>CUDAGraphRunner</code> 只用在了下面两个地方。</p>
<h3>ktransformers/server/backend/interfaces/ktransformers.py</h3>
<p>上文提到的 <code>KTransformersInterface::decode_one_tokens()</code> 实现,重载了 <code>TransformersInterface</code> 对应方法。<br>
cuda graph 相关的实现也比较简单。如果是首次调用该函数会初始化 <code>CUDAGraphRunner</code>,然后 <code>CUDAGraphRunner::capture</code> 启动图捕获,借助 model 进行推理得到 logits 转换成 token。<br>
然后再次调用时发现 <code>CUDAGraphRunner</code> 已经构造好了,就直接把参数传给 <code>CUDAGraphRunner::forward</code>。</p>
<h3>ktransformers/util/utils.py</h3>
<p><code>prefill_and_generate</code>,这个函数就是给 <code>ktransformers/local_chat.py</code> 用的,它是个用来调试的命令行工具。<br>
值得注意的是整个项目目前支持的模型列表也是定义在 <code>ktransformers/local_chat.py</code> 这里的 = = 这个结构组织得有点乱,于是你可以看到在 KTransformersInterface 里面是 <code>from ktransformers.local_chat import custom_models, default_optimize_rules</code> 获取支持的模型列表和用于调优的 YAML 文件。</p>
<pre><code class="language-python">custom_models = {
"DeepseekV2ForCausalLM": DeepseekV2ForCausalLM,
"Qwen2MoeForCausalLM": Qwen2MoeForCausalLM,
"MixtralForCausalLM": MixtralForCausalLM,
}
ktransformer_rules_dir = os.path.dirname(os.path.abspath(__file__)) + "/optimize/optimize_rules/"
default_optimize_rules ={
"DeepseekV2ForCausalLM": ktransformer_rules_dir + "DeepSeek-V2-Chat.yaml",
"Qwen2MoeForCausalLM": ktransformer_rules_dir + "Qwen2-57B-A14B-Instruct.yaml",
"MixtralForCausalLM": ktransformer_rules_dir + "Mixtral.yaml",
}
</code></pre>
]]></content:encoded>
<author>peihao.young@gmail.com (Peihao Yang)</author>
</item>
<item>
<title><![CDATA[TensorRT-LLM 源码阅读]]></title>
<link>https://forsworns.github.io///zh/blogs/20240916/</link>
<guid>https://forsworns.github.io///zh/blogs/20240916/</guid>
<pubDate>Mon, 16 Sep 2024 00:00:00 GMT</pubDate>
<description><![CDATA[再看看 cuda graph 咋开启的 ...]]></description>
<content:encoded><![CDATA[<p>[[toc]]</p>
<p>基于 TensorRT-LLM-0.10.0,<a href="https://github.com/NVIDIA/TensorRT-LLM/releases/tag/v0.10.0">https://github.com/NVIDIA/TensorRT-LLM/releases/tag/v0.10.0</a><br>
最开始想看下 cuda graph 怎么开启的。<br>
trtllm-build 工具不像 trtexec,有 trtexec --useCudaGraph 这个选项</p>
<h1>cpp/tensorrt_llm/runtime/gptSession.cpp</h1>
<p>cpp runtime 中关于 cuda graph API,是 <code>GptSession::mCudaGraphMode</code> 这个变量控制的,它被设置成了外部的 <code>tr::GptSession::Config::cudaGraphMode</code> 的配置值。</p>
<h1>cpp/tensorrt_llm/pybind/bindings.cpp</h1>
<p>python binding 文件,可以看到例如 cpp 中的 <code>tr::GptSession::Config</code> 被映射成了 python 中的 <code>GptSessionConfig</code>。<br>
它的成员被映射成了 <code>GptSessionConfig::cuda_graph_mode</code>。</p>
<h1>tensorrt_llm/runtime/model_runner_cpp.py</h1>
<p>包装了 cpp 目录下的具体实现,暴露成 python 接口。从 <code>ModelRunnerCppGptSession::from_dir()</code> 的实现里面实际上可以看到 <code>GptSessionConfig</code> 这些配置项都是怎么传递进去的。实际上没有传递 cuda graph 那个参数。</p>
<h1>tensorrt_llm/runtime/generation.py</h1>
<p>python runtime 中,同样是在 decode 阶段使用 cuda graph 加速。</p>
<p>在调用 <code>tensorrt_llm.runtime.GenerationSession()</code> 的时候,配置一下 <code>cuda_graph_mode=True</code> 即可。也就是改一下例如 <code>/examples/llama/summarize_long.py</code> 这样的示例代码。<br>
tensorrt-llm 自己的 benchmark 里面倒是加了对应的开关,只不过是示例代码里面省略了。</p>
]]></content:encoded>
<author>peihao.young@gmail.com (Peihao Yang)</author>
</item>
<item>
<title><![CDATA[ollama 源码阅读]]></title>
<link>https://forsworns.github.io///zh/blogs/20240909/</link>
<guid>https://forsworns.github.io///zh/blogs/20240909/</guid>
<pubDate>Mon, 09 Sep 2024 00:00:00 GMT</pubDate>
<description><![CDATA[推理服务器]]></description>
<content:encoded><![CDATA[<p>[[toc]]</p>
<p>基于 <a href="https://github.com/ollama/ollama/tree/123a722a6f541e300bc8e34297ac378ebe23f527">https://github.com/ollama/ollama/tree/123a722a6f541e300bc8e34297ac378ebe23f527</a></p>
<p>ollama 是一个通用的 llm 推理服务器,借助 llama.cpp 进行推理。<br>
ollama 0.1.44 镜像内,将 llama.cpp 编写的推理服务器放到了 /tmp/ollama1917690259/runners/下,以 cuda 后端为例,为 /tmp/ollama1917690259/runners/cuda_v11/ollama_llama_server。<br>
当我们调用 ollama serve后,只会启动一个 go 编写的 server;当我们执行 ollama run qwen:7b它会拉起 llama.cpp server,然后作为反向代理转发我们的请求给 lamma.cpp server。</p>
<h2>gpu/gpu.go</h2>
<p>一开始看 ollama 的代码是碰到了问题,想看下 go server 侧为什么会调用到 cuda API。</p>
<p>GetGPUInfo()->initCudaHandles()这里会打开 <a href="http://libnvidia-ml.so">libnvidia-ml.so</a>、<a href="http://libcuda.so">libcuda.so</a>、<a href="http://libcudart.so">libcudart.so</a>。<br>
会去找这些库中的一些符号,go server 运行期间通过 cgo 调用他们。</p>
<h2>server/routes.go</h2>
<p><code>func Serve(ln net.Listener)</code> 函数也就是调用 ollama serve时执行的函数,它在启动 go server 前会去调用上面提到的<code>GetGPUInfo()</code>获取 GPU 信息。</p>
<p>看下这个文件的其他内容</p>
<p><code>func (s *Server) GenerateRoutes()</code> 配置路由和对应的 handler。<br>
以 <code>r.POST("/api/embeddings", s.EmbeddingsHandler)</code> 为例,<code>func (s *Server) EmbeddingsHandler(c *gin.Context)</code> 中解析请求参数,调用 <code>s.sched.GetRunner</code> 阻塞直到获取到 runner。<br>
再调用 runner.llama.Embedding 获取 llama.cpp 中的模型服务,获取响应返回给用户。</p>
<pre><code class="language-go">func (s *Server) EmbeddingsHandler(c *gin.Context) {
var req api.EmbeddingRequest
err := c.ShouldBindJSON(&req)
model, err := GetModel(req.Model)
opts, err := modelOptions(model, req.Options)
rCh, eCh := s.sched.GetRunner(c.Request.Context(), model, opts, req.KeepAlive.Duration)
var runner *runnerRef
select {
case runner = <-rCh:
case err = <-eCh:
handleErrorResponse(c, err)
return
}
embedding, err := runner.llama.Embedding(c.Request.Context(), req.Prompt)
resp := api.EmbeddingResponse{
Embedding: embedding,
}
c.JSON(http.StatusOK, resp)
}
</code></pre>
<h2>server/sched.go</h2>
<p>除了直接调用 <code>GetGPUInfo()</code>,该函数还可能通过该文件下的 <code>Scheduler.getGpuFn</code> 函数指针调用,包含下面两个调用处</p>
<ul>
<li>func (s *Scheduler) processPending(ctx context.Context)</li>
<li>func (runner *runnerRef) waitForVRAMRecovery()</li>
</ul>
<p>看下这个文件的其他内容</p>
<p><code>runnerRef</code> 是调度的实体,对应请求中的 <code>req.model.ModelPath</code>,为这个模型启动 llama.cpp 服务器。</p>
<p><code>GetRunner</code>,把用户请求 req 写入了 <code>s.pendingReqCh</code>,如果失败了把错误写入到 <code>req.errCh</code>,没有失败的时候,会阻塞在 <code>req.successCh</code>。</p>
<pre><code class="language-go">func (s *Scheduler) GetRunner(c context.Context, model *Model, opts api.Options, sessionDuration time.Duration) (chan *runnerRef, chan error) {
req := &LlmRequest{
ctx: c,
model: model,
opts: opts,
sessionDuration: sessionDuration,
successCh: make(chan *runnerRef),
errCh: make(chan error, 1),
}
select {
case s.pendingReqCh <- req:
default:
req.errCh <- ErrMaxQueue
}
return req.successCh, req.errCh
}
</code></pre>
<p><code>Run</code> 函数会创建两个 go routine 去分别处理等待队列和完成队列,刚刚的 <code>s.pendingReqCh </code> 中的请求就是在 <code>processPending</code> 中处理的。任务成功后写回到 <code>req.successCh</code>。</p>
<pre><code class="language-go">func (s *Scheduler) Run(ctx context.Context) {
go func() {
s.processPending(ctx)
}()
go func() {
s.processCompleted(ctx)
}()
}
</code></pre>
<p><code>func (s *Scheduler) load(req *LlmRequest, ggml *llm.GGML, gpus gpu.GpuInfoList)</code> 调用 <code>Scheduler.newServerFn</code>,也就是 <code>llm/server.go</code> 中的 <code>NewLlamaServer</code> 创建 llama.cpp 服务;同时创建调度实体 <code>runnerRef</code>。</p>
<p><code>func (s *Scheduler) findRunnerToUnload()</code> 用来寻找一个最合适被关闭的 runnerRef,会先去看 runner.refCount 这个引用计数,看是否有空闲的 runnerRef,如果有就把它关闭;否则就给所有 runnerRef 按 runnerRef.sessionDuration 排序,返回马上要执行完成的 runner。</p>
<h2>initCudaHandles中寻找的符号</h2>
<ul>
<li>gpu/gpu_info_nvcuda.h:cuda_handle_t</li>
<li>gpu/gpu_info_cudart.h:cudart_handle_t</li>
<li>gpu/gpu_info_nvml.h:nvml_handle_t</li>
</ul>
]]></content:encoded>
<author>peihao.young@gmail.com (Peihao Yang)</author>
</item>
<item>
<title><![CDATA[Let's reproduce GPT-2 笔记]]></title>
<link>https://forsworns.github.io///zh/blogs/20240626/</link>
<guid>https://forsworns.github.io///zh/blogs/20240626/</guid>
<pubDate>Wed, 26 Jun 2024 00:00:00 GMT</pubDate>
<description><![CDATA[学习了一下 Andrej Karpathy 大神的 GPT-2 视频课程]]></description>
<content:encoded><![CDATA[<p>[[toc]]</p>
<p>之前看了 Andrej Karpathy 的 Tokenizer 视频,最近他又发了一个从零复现 GPT-2 的视频,学习一个。</p>
<p>相关博客</p>
<p>网上读到 Maeiee 记录的,Karpathy 之前另一期从零搭建 GPT 视频的学习笔记 <a href="https://garden.maxieewong.com/087.%E8%A7%86%E9%A2%91%E5%BA%93/YouTube/Andrej%20Karpathy/Let%27s%20build%20GPT%EF%BC%9Afrom%20scratch,%20in%20code,%20spelled%20out./#step6-bigram-language-model-v1">Let's build GPT:from scratch, in code, spelled out.</a>。</p>
<p>仓库地址</p>
<ul>
<li>视频配套的教学仓库 build-nanogpt: <a href="https://github.com/karpathy/build-nanogpt">https://github.com/karpathy/build-nanogpt</a></li>
<li>nanoGPT: <a href="https://github.com/karpathy/nanoGPT">https://github.com/karpathy/nanoGPT</a></li>
<li>llm.c: <a href="https://github.com/karpathy/llm.c">https://github.com/karpathy/llm.c</a></li>
</ul>
<p><a href="https://www.youtube.com/watch?v=l8pRSuU81PU">视频地址</a></p>
<p>论文</p>
<ul>
<li>Transformer:<a href="https://arxiv.org/abs/1706.03762">Attention is All You Need</a></li>
<li>GPT-2:<a href="https://d4mucfpksywv.cloudfront.net/better-language-models/language_models_are_unsupervised_multitask_learners.pdf">Language Models are Unsupervised Multitask Learners</a>,论文中的参数配置信息不够具体,需要参考 GPT-3 的论文。</li>
<li>GPT-3:<a href="https://arxiv.org/abs/2005.14165">Language Models are Few-Shot Learners</a></li>
</ul>
<h2>引言</h2>
<p>复现的是 124M 的模型,原始论文中的参数统计数据有误。该模型由 12个 768 channel、768 dimension 的 transformer 构成。</p>
<p>借助 HuggingFace 的 transformers 库,打印了 GPT-2 中的张量信息。首先是输入层,tokenizer embedding 的规模是 50257,每个 token 被表示为 768 维的向量,所以 wte 是一个 50257x768 的矩阵;position embedding 上下文长度为 1024,每个位置由 768 维的向量编码,所以 wpe 是一个 1024x768 的矩阵。<br>
原始 transformer 文章使用固定的正弦余弦波来做位置编码(考虑到正弦波的加法性质),GPT-2将位置编码也参数化了。</p>
]]></content:encoded>
<author>peihao.young@gmail.com (Peihao Yang)</author>
</item>
<item>
<title><![CDATA[ollama/llama.cpp 源码阅读]]></title>
<link>https://forsworns.github.io///zh/blogs/20240623/</link>
<guid>https://forsworns.github.io///zh/blogs/20240623/</guid>
<pubDate>Sun, 23 Jun 2024 00:00:00 GMT</pubDate>
<description><![CDATA[最开始是想看下为什么 cuda graph 没有被启用]]></description>
<content:encoded><![CDATA[<p>[[toc]]</p>
<h2>ollama</h2>
<p>ollama基于 <a href="https://github.com/ollama/ollama/blob/ccef9431c8aae4ecfd0eec6e10377d09cb42f634">https://github.com/ollama/ollama/blob/ccef9431c8aae4ecfd0eec6e10377d09cb42f634</a></p>
<h3>llm/</h3>
<h4>server.go</h4>
<p>go 写的 server,主体为 <a href="https://github.com/ollama/ollama/blob/ccef9431c8aae4ecfd0eec6e10377d09cb42f634/llm/server.go#L80"><code>NewLlamaServer</code></a></p>
<p>它会拉起多个进程,分别执行下面的 ext_server/server.cpp 中,基于 llama.cpp 实现的真正做推理服务的 server。</p>
<h4>ext_server/server.cpp</h4>
<p>把 llama.cpp 导入成了一个 submodule,基于 llama.cpp 开发的一个推理服务器。</p>
<h2>llama.cpp</h2>
<p>llama.cpp 基于 <a href="https://github1s.com/ggerganov/llama.cpp/blob/45c0e2e4c1268c2d7c8c45536f15e3c9a731ecdc/llama.h">https://github1s.com/ggerganov/llama.cpp/blob/45c0e2e4c1268c2d7c8c45536f15e3c9a731ecdc/llama.h</a></p>
<p>看编译脚本上是默认开 cuda graph 优化,但是用 ollama 起的服务器跑的时候没有用到。</p>
<p>CmakeLists.txt 里面声明了只要找到了 libcuda,就会定义 GGML_CUDA_USE_GRAPHS 开启 cuda graph 优化。Makefile 中同样,只要是声明了 <code>make LLAMA_CUDA=1</code> 就会定义 <code>GGML_CUDA_USE_GRAPHS</code> 开启 cuda graph 优化。</p>
<h3>llama.cpp</h3>
<p>对外的 llama.cpp 库 API 实现</p>
<h3>ggml.c</h3>
<p>被 llama.cpp 包了一层的内部 API</p>
<h3>ggml-backend.c</h3>
<p>不同的后端通过 <code>ggml_backend_register</code> 注册自身,<code>ggml_backend_registry_init</code> 运行时分别调用他们,这里利用了一个技巧避免引入头文件。</p>
<pre><code class="language-cpp">GGML_CALL static void ggml_backend_registry_init(void) {
ggml_backend_register("CPU", ggml_backend_reg_cpu_init, ggml_backend_cpu_buffer_type(), NULL);
// add forward decls here to avoid including the backend headers
#ifdef GGML_USE_CUDA
extern GGML_CALL void ggml_backend_cuda_reg_devices(void);
ggml_backend_cuda_reg_devices();
#endif
// …
}
</code></pre>
<p>实现 <code>static struct ggml_backend_i cpu_backend_i </code> 后端。</p>
<h3>ggml-blas.cpp</h3>
<p>实现 <code>static struct ggml_backend_i blas_backend_i</code> 后端。</p>
<h3><a href="http://ggml-cuda.cu">ggml-cuda.cu</a></h3>
<p>实现 <code>static ggml_backend_i ggml_backend_cuda_interface</code> 后端。</p>
<p>cuda graph 是由 <a href="https://github.com/ggerganov/llama.cpp/commit/bc4bba364fb96d908f2698e908648df5e6f55e02">https://github.com/ggerganov/llama.cpp/commit/bc4bba364fb96d908f2698e908648df5e6f55e02</a> 这个 commit-bc4b 引入的。</p>
<h3>ggml-cuda</h3>
<h4>cpy.cuh</h4>
<p>commit-bc4b 为 <code>struct ggml_backend_cuda_context</code> 新增了一个成员,<code>std::unique_ptr<ggml_cuda_graph> cuda_graph</code>。看上去一个 context 只会捕获出一个 cuda graph。</p>
<p>结构体 <code>ggml_cuda_graph</code> 在析构的时候会自动调用 <code>cudaGraphExecDestroy</code> 和 <code>cudaGraphDestroy</code> 清理之前捕获到的 cuda graph。它的定义比较简单,如下</p>
<pre><code class="language-cpp">struct ggml_cuda_graph {
cudaGraph_t graph = nullptr;
cudaGraphExec_t instance = nullptr;
size_t num_nodes = 0;
std::vector<cudaGraphNode_t> nodes;
std::vector<cudaKernelNodeParams> params;
// 禁用该 feature 的几种可能的原因
bool disable_due_to_gpu_arch = false;
// 如果当前用例中,图节点更新得太快,那图需要一直重建,建图的开销可能会大于 cuda graph 节省的开销。
bool disable_due_to_too_many_updates = false;
bool disable_due_to_failed_graph_capture = false;
int number_consecutive_updates = 0;
std::vector<ggml_graph_node_properties> ggml_graph_properties;
std::vector<char **> updated_kernel_arg;
};
struct ggml_graph_node_properties {
void * node_address;
// 同 `ggml_tensor` 上的 `ggml_op`,例如 `GGML_OP_CPY`、`GGML_OP_VIEW`
ggml_op node_op;
// 同 `ggml_tensor` 上的 `ne`
int64_t ne[GGML_MAX_DIMS];
// 同 `ggml_tensor` 上的 `nb`
size_t nb[GGML_MAX_DIMS];
// 同 `ggml_tensor` 上的 `src[i]->data`
void * src_address[GGML_MAX_SRC];
};
</code></pre>
<p><code>set_ggml_graph_node_properties</code> 从一个 <code>ggml_tensor</code> 构建一个 <code>ggml_graph_node_properties</code>,转换成图里的节点。<br>
<code>ggml_graph_node_has_matching_properties</code> 比较 <code>ggml_tensor</code> 和 <code>ggml_graph_node_properties</code> 的成员,判断二者是否匹配。</p>
<p><code>ggml_backend_cuda_graph_compute</code> 中根据参数 <code>ggml_backend_t</code> 和 <code>ggml_cgraph</code> 去构建 <code>ggml_backend_t->ggml_backend_cuda_context->ggml_cuda_graph</code>。这个函数里面首先检查了是不是安培以下的 GPU,如果是,就不用 cuda graph 了。之前在 T4 上测试的,所以没用到 cuda graph。有点坑爹 release 模式下不开 <code>LLAMA_DEBUG</code>,这错误日志就不打了。</p>
<pre><code class="language-cpp">GGML_CALL static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend, ggml_cgraph * cgraph) {
// ...
if (cuda_ctx->cuda_graph->graph == nullptr) {
if (ggml_cuda_info().devices[cuda_ctx->device].cc < CC_AMPERE) {
cuda_ctx->cuda_graph->disable_due_to_gpu_arch = true;
#ifndef NDEBUG
GGML_CUDA_LOG_WARN("%s: disabling CUDA graphs due to GPU architecture\n", __func__);
#endif
}
}
// ...
}
</code></pre>
<p>如果启用 cuda graph,则比较当前传入的 <code>ggml_cgraph</code> 和之前当前的 cuda graph 是否相同</p>
<pre><code class="language-cpp">GGML_CALL static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend, ggml_cgraph * cgraph) {
// ...
if (cuda_ctx->cuda_graph->instance == nullptr) {
cuda_graph_update_required = true;
}
// Check if the graph size has changed
if (cuda_ctx->cuda_graph->ggml_graph_properties.size() != (size_t)cgraph->n_nodes) {
cuda_graph_update_required = true;
cuda_ctx->cuda_graph->ggml_graph_properties.resize(cgraph->n_nodes);
}
// Loop over nodes in GGML graph to determine if CUDA graph update is required
// and store properties to allow this comparison for the next token
for (int i = 0; i < cgraph->n_nodes; i++) {
bool has_matching_properties = true;
if (!cuda_graph_update_required) {
has_matching_properties = ggml_graph_node_has_matching_properties(cgraph->nodes[i], &cuda_ctx->cuda_graph->ggml_graph_properties[i]);
}
if (!has_matching_properties) {
cuda_graph_update_required = true;
}
set_ggml_graph_node_properties(cgraph->nodes[i], &cuda_ctx->cuda_graph->ggml_graph_properties[i]);
}
// ...
}
</code></pre>
<p>再次遍历当前的 <code>ggml_cgraph</code>,更新 <code>GGML_OP_CPY</code> 类型节点的信息,因为拷贝操作的地址会随着 token 变化。</p>
<pre><code class="language-cpp">GGML_CALL static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend, ggml_cgraph * cgraph) {
// ...
// Loop over nodes in GGML graph to obtain info needed for CUDA graph
cuda_ctx->cuda_graph->updated_kernel_arg.clear();
for (int i = 0; i < cgraph->n_nodes; i++) {
ggml_tensor * node = cgraph->nodes[i];
if (node->src[0] && ggml_backend_buffer_is_cuda_split(node->src[0]->buffer)) {
use_cuda_graph = false; // Split buffers are not supported by CUDA graph capture
}
if (node->op == GGML_OP_MUL_MAT_ID) {
use_cuda_graph = false; // This node type is not supported by CUDA graph capture
}
if (node->op == GGML_OP_ADD && node->src[1] && node->src[1]->ne[1] > 1) {
// disable CUDA graphs for batch size > 1 for now.
// Changes in batch size or context size can cause changes to the grid size of some kernels.
use_cuda_graph = false;
}
if (node->op == GGML_OP_CPY) {
// store the copy op parameter which changes with each token.
cuda_ctx->cuda_graph->updated_kernel_arg.push_back((char **) &(node->src[1]->data));
// store a pointer to each copy op CUDA kernel to identify it later
void * ptr = ggml_cuda_cpy_fn(node->src[0], node->src[1]);
if (std::find(ggml_cuda_cpy_fn_ptrs.begin(), ggml_cuda_cpy_fn_ptrs.end(), ptr) == ggml_cuda_cpy_fn_ptrs.end()) {
ggml_cuda_cpy_fn_ptrs.push_back(ptr);
}
}