לדלג לתוכן

6.2 cuDNN primitives לרשתות עמוקות הרצאה

בשיעור 6.1 הכרנו את cuBLAS - הספרייה שנותנת לנו kernel מהיר וסגור לאלגברה לינארית כללית, בוחרת בעצמה את המימוש הטוב ביותר לפי הוריסטיקה אטומה, ומצפה למטריצות בסדר column-major. cuBLAS מכסה יפה את הכפל הצפוף (dense GEMM) שיושב בלב שכבות ה-linear של רשת עצבית, אבל רשת עמוקה מודרנית עשויה מהרבה יותר מכפל מטריצות: קונבולוציות (convolutions), נרמולים (normalizations), pooling, פונקציות אקטיבציה, ומנגנוני attention. עבור אלה NVIDIA מספקת ספרייה אחות ל-cuBLAS: cuDNN - ספריית רשתות עצביות עמוקות ל-CUDA (CUDA Deep Neural Network). בשיעור הזה נבין מה cuDNN מכילה, למה ה-Graph API החדש שלה הוא קפיצת מדרגה על פני הקריאה הישירה לפעולה בודדת, איך fusion של פעולות לתוך kernel יחיד חוסך רוחב פס לזיכרון, ולמה - למרות כל זה - מרכז הכובד של cuDNN נשאר ברשתות קונבולוציה על ארכיטקטורת Ampere, בעוד שעבור Transformers על Hopper ו-Blackwell מפנה NVIDIA את המפתחים אל CUTLASS, נושא שיפתח את שיעור 6.3.

מהי cuDNN - CUDA Deep Neural Network

cuDNN היא, בלשון ה-glossary, "ספריית primitives לבניית רשתות עצביות עמוקות מואצות-GPU". המילה primitives היא המפתח: אלה אבני הבניין הבסיסיות שכל רשת עמוקה מורכבת מהן, כתובות פעם אחת, ומכוונות ביד עד קצה היכולת של החומרה, כדי שכותבי ה-frameworks (כמו PyTorch ו-TensorFlow) לא ייאלצו לממש אותן מחדש בעצמם. בדיוק כפי ש-cuBLAS חוסך מאיתנו לכתוב GEMM מאופס, cuDNN חוסך מאיתנו לכתוב קונבולציה, batch normalization או attention מאופסים.

חשוב למקם אותה נכון במחסנית התוכנה שמיפינו בפרק 5. cuDNN יושבת בשכבת האפליקציה של פלטפורמת CUDA, בדיוק לצד cuBLAS, מעל ה-Runtime API ומעל הדרייבר. היא איננה חלק מהדרייבר ואיננה חלק מ-CUDA Toolkit הבסיסי - היא חבילה נפרדת שמתקינים בנפרד (וה-frameworks מביאים איתם גרסה משלהם).

   cuDNN's place in the software stack (continuation of chapter 5):

   ┌─────────────────────────────────────────────┐
   │  framework:  PyTorch / TensorFlow / JAX      │
   ├──────────────────────┬──────────────────────┤
   │  cuBLAS              │  cuDNN               │  ← the libraries layer (6.1, 6.2)
   │  (GEMM, general algebra) │  (conv, norm, attn)  │
   ├──────────────────────┴──────────────────────┤
   │  CUDA Runtime API  (libcudart)  ← chapter 5  │
   ├─────────────────────────────────────────────┤
   │  Driver API / libcuda / nvidia.ko            │
   ├─────────────────────────────────────────────┤
   │  GPU hardware  (H100: 132 SMs, Tensor Cores) │  ← chapter 1
   └─────────────────────────────────────────────┘

חלוקת העבודה בין שתי האחיות ברורה, ושווה לזכור אותה: ה-frameworks נשענים על cuBLAS לאלגברה לינארית כללית - הכפלים הצפופים של שכבות ה-linear - ועל cuDNN ל-primitives המתמחים: שכבות קונבולוציה, נרמולים, ו-attention. שתיהן סגורות, שתיהן הוריסטיות, ושתיהן משרתות את אותו לקוח - ה-framework שמעליהן.

קטלוג ה-primitives - מקונבולוציה ל-attention

בניגוד לרושם הרווח, cuDNN כבר מזמן אינה "ספריית הקונבולוציות". נכון שהקונבולוציה היא הליבה ההיסטורית שלה - cuDNN נולדה בעידן שבו רשתות קונבולוציה (CNNs) שלטו בראייה ממוחשבת, וזו עדיין נקודת החוזק המזוהה איתה ביותר - אבל הקטלוג של היום רחב בהרבה. הטבלה הבאה מסכמת את משפחות ה-primitives המרכזיות ואת מקומן ברשת:

primitive מה זה עושה היכן ברשת
Convolution קונבולוציה דו/תלת-ממדית, ליבת cuDNN ההיסטורית שכבות conv ב-CNN (ResNet, U-Net)
Attention / scaled dot-product attention בסגנון FlashAttention, כולל softmax מובנה בלוקי ה-Transformer
Matrix multiplication כפל מטריצות, לרוב בשילוב עם epilogue מותך linear / projection layers
Normalizations batch/layer/group normalization אחרי conv או בתוך בלוק Transformer
Pooling max/average pooling, דגימה-למטה בין שכבות conv ב-CNN
Activations ReLU, GELU, sigmoid וכו' אחרי כמעט כל שכבה

שני פריטים בטבלה ראויים להדגשה. הראשון הוא ה-attention: cuDNN מספקת מימוש של scaled dot-product attention בסגנון FlashAttention - כלומר attention שמחשב את softmax בזרימה (streaming) מבלי לכתוב את מטריצת ה-scores המלאה בגודל N×N לזיכרון הגלובלי. זה הופך את cuDNN מספריית CNN טהורה לספרייה שרלוונטית גם ל-Transformers, לפחות ברמת ה-primitive. השני הוא ה-matmul: cuDNN אכן יודעת לכפול מטריצות, אבל לא כתחליף ל-cuBLAS - היא עושה זאת בעיקר כדי לשלב את הכפל בתוך גרף פעולות מותך, נושא הסעיף הבא.

ה-Graph API - תיאור דקלרטיבי של גרף פעולות

הדרך המודרנית לעבוד עם cuDNN אינה קריאה לפונקציה בודדת פר פעולה, אלא בנייה של גרף פעולות (operation graph) דרך ה-Graph API הדקלרטיבי. הרעיון קרוב מאוד למה שראינו בשיעור 5.5 על גרפי CUDA, אבל ברמת הפשטה גבוהה יותר: שם תיארנו גרף של kernel launches, וכאן אנחנו מתארים גרף של פעולות מתמטיות מופשטות (conv, bias, activation) ונותנים ל-cuDNN לבחור אילו kernels יממשו אותן.

ההבדל בין דקלרטיבי לאימפרטיבי הוא לב העניין. בסגנון האימפרטיבי הישן היינו קוראים לפונקציה, מקבלים תוצאה בזיכרון הגלובלי, קוראים לפונקציה הבאה עליה, וכן הלאה - כל פעולה אטומה וחסרת מודעות לשכנותיה. בסגנון הדקלרטיבי אנחנו מתארים ל-cuDNN את כל רצף הפעולות מראש, כטופולוגיה, ורק אז מבקשים ממנה לממש אותו. משום שהיא רואה את הגרף השלם לפני שהיא מהדרת משהו, היא יכולה לנתח אותו ולבצע אופטימיזציות חוצות-פעולה שאי אפשר לבצע כשרואים פעולה אחת בכל פעם - בראש ובראשונה fusion, שנראה בסעיף הבא.

מבחינת ה-API עצמו, יש כאן מבנה של frontend פתוח מעל backend סגור, וכדאי להבין אותו לעומק:

   ┌──────────────────────────────────────────────┐
   │  cuDNN frontend  (open-source)                │
   │  Python  +  C++                               │
   │  this is where the operation graph is built   │
   │  declaratively                                │
   └───────────────────┬──────────────────────────┘
                       │  calls into
   ┌───────────────────▼──────────────────────────┐
   │  cuDNN C backend  (closed-source)             │
   │  analyzes the graph, chooses implementations,  │
   │  compiles kernels                             │
   └──────────────────────────────────────────────┘

השכבה העליונה, ה-frontend, היא ספרייה בקוד פתוח עם interfaces Python ו-C++. זו השכבה שבה כותבים בפועל - מגדירים tensors, מגדירים פעולות, ומחברים אותן לגרף. השכבה התחתונה, ה-backend, היא ספריית C בקוד סגור. היא זו שמקבלת את הגרף המתואר, מנתחת אותו, בוחרת את המימושים הקונקרטיים ומהדרת אותם. ה-backend הסגור נגיש גם ישירות דרך C FFI עבור מקרי שימוש legacy, אבל המפתח המודרני עובד כמעט תמיד דרך ה-frontend הפתוח.

בקוד Python מינימלי, בניית גרף conv-bias-relu נראית עקרונית כך (סכמטי, להמחשת הזרימה):

import cudnn

graph = cudnn.pygraph(io_data_type=cudnn.data_type.HALF,
                      compute_data_type=cudnn.data_type.FLOAT)

X = graph.tensor(name="X", dim=[N, C, H, W])      # input
W_filt = graph.tensor(name="W", dim=[K, C, R, S])  # convolution filter
B = graph.tensor(name="B", dim=[1, K, 1, 1])       # bias

# describing the graph declaratively, operation by operation:
conv_out = graph.conv_fprop(X, W_filt, padding=[1, 1], stride=[1, 1])
bias_out = graph.bias(conv_out, B)
Y = graph.relu(bias_out)
Y.set_output(True)

# only now does cuDNN see the whole graph and build an execution plan:
graph.build([cudnn.heur_mode.A])

שימו לב לסדר: כל שורות ה-tensor וה-conv/bias/relu הן תיאור בלבד - הן לא מריצות דבר על ה-GPU. רק הקריאה ל-build היא זו ש-cuDNN מנתחת בה את הגרף השלם, מחליטה על fusion, בוחרת מימושים, ומכינה את התוכנית. זה בדיוק הפיצול תיאור/הכנה/הרצה שראינו בגרפי CUDA בשיעור 5.5, ומאותו טעם: את העבודה היקרה של הניתוח והבחירה עושים פעם אחת, ומופחתים (amortize) אותה על פני אלפי הרצות.

מיזוג פעולות - Fusion - kernel אחד במקום שלושה

ה-fusion (מיזוג פעולות) הוא הרווח המרכזי שה-Graph API מאפשר, ושווה להבין אותו דרך המספרים, כי הוא סיפור על רוחב פס לזיכרון ולא על כמות חישוב. ניקח את הרצף הקלאסי Convolution + Bias + ReLU, שמופיע פעמים אינספור ברשת קונבולוציה.

בלי fusion, כל פעולה היא kernel נפרד, וכל kernel חייב לקרוא את הקלט שלו מהזיכרון הגלובלי (global memory, ה-VRAM של הכרטיס) ולכתוב את הפלט שלו בחזרה אליו:

   without fusion (three kernels, three round trips to global memory):

   [conv] ── writes ──► intermediate tensor (global) ── reads ──► [bias]
                        │                                    │
                        └─ writes ──► intermediate tensor ─ reads ─► [relu] ──► output

התוצאה של ה-conv נכתבת במלואה ל-global memory, נקראת ממנו במלואה על ידי ה-bias, נכתבת שוב, ונקראת שוב על ידי ה-relu. עבור tensor ביניים גדול, אלה מגה-בייטים רבים שנוסעים הלוך-ושוב על ערוץ ה-HBM3 - וה-bias וה-relu הם פעולות אלמנט-אחר-אלמנט זולות מאוד חישובית, כך ששני ה-kernels האלה memory-bound לחלוטין. אנחנו משלמים את מלוא מחיר רוחב הפס על חישוב פעוט.

עם fusion, cuDNN ממזגת את השלושה ל-kernel יחיד. תוצאת הביניים של ה-conv נשמרת ב-shared memory (זיכרון משותף על ה-SM, מהיר בסדר גודל מה-global) או באוגרים, וה-bias וה-relu מוחלים עליה שם, בלי לגעת בזיכרון הגלובלי כלל:

   with fusion (one kernel, one round trip to global memory):

   [conv → bias → relu]   ← the intermediate lives in shared memory / registers
        │
        └── writes once ──► output (global)

נעשה חשבון גס על tensor ביניים ב-FP16 בגודל N=32, C=256, H=W=56, כלומר כ-32·256·56·56 ≈ 25.7 מיליון אלמנטים, בערך 51 מגה-בייט:

   global memory traffic for the bias+relu stage:

   without fusion:  read 51MB + write 51MB (bias)  +
                    read 51MB + write 51MB (relu)   =  ~204 MB

   with fusion:      the intermediate never goes to global at all  =  ~0 MB extra

   on the H100's HBM3 (~3.35 TB/s), a saving of ~204MB
   is equivalent to a saving of ~61 microseconds of memory time - and over
   millions of calls across a whole network, this becomes significant.

הרווח, אם כן, אינו בכמות ה-FLOPs - אותו מספר כפל-חיבור מתבצע כך או כך - אלא בכך שאנחנו מוחקים את נסיעות הזיכרון של תוצאות הביניים. זה בדיוק ההיגיון של ה-tiling ב-shared memory שראינו בשיעור 3.3, מוחל כאן אוטומטית על ידי הספרייה במקום ידנית על ידינו. וזו גם הסיבה שה-fusion חייב את הראייה הגלובלית של הגרף: kernel של conv בודד לא יודע ש-bias ו-relu עומדים לבוא אחריו, ולכן חייב לכתוב את תוצאתו החוצה; רק כש-cuDNN רואה את שלושתם יחד היא יכולה להחליט להשאיר את הביניים על ה-chip.

בחירת מימוש הוריסטית - אותה פילוסופיה של cuBLAS

כמו cuBLAS, גם cuDNN אינה מריצה "מימוש אחד" לכל פעולה. לכל primitive היא מחזיקה מספר מימושים בסיסיים שונים, ובוחרת ביניהם בזמן ריצה באמצעות הוריסטיקה פנימית - שה-glossary מסמן במפורש, שוב, כ"(unknown)", כלומר אטומה ולא מתועדת. הבחירה נעשית לפי שלושה צירים:

  • ארכיטקטורת ה-SM של הכרטיס (Ampere מול Hopper מול Blackwell) - קונבולוציה על H100 עשויה להיבחר אחרת מאשר על A100.
  • ה-data type של הקלט (FP32, FP16, BF16, INT8) - מימוש שמנצל Tensor Cores ב-FP16 שונה מזה שרץ על CUDA Cores ב-FP32.
  • גודל הקלט - ה-batch, מספר הערוצים, ממדי הפילטר וה-spatial. קונבולוציה 1×1 על tensor רזה נבחרת אחרת מקונבולוציה 3×3 על tensor שמן.

זו בדיוק אותה פילוסופיה שראינו ב-6.1: הרבה kernels קונקרטיים לכל פעולה לוגית אחת, והמפתח אינו בוחר ביניהם - הספרייה בוחרת. אתם מתארים את הגרף באופן דקלרטיבי, ו-cuDNN בעלת הבעלות על הבחירה של המימוש הקונקרטי. היתרון ברור - אתם מקבלים kernel קרוב-לאופטימלי בלי להיות מומחי חומרה - אבל יש לזה מחיר ותופעות לוואי שכדאי להכיר:

  • ביצועים לא צפויים לחלוטין. אותה שכבה עם batch מעט שונה או dtype שונה עלולה ליפול על מימוש אחר לגמרי, עם ביצועים שונים מהותית. שינוי קטן בקלט יכול לשנות את הבחירה.
  • תלות בגרסה ובארכיטקטורה. מעבר מ-A100 ל-H100, או שדרוג גרסת cuDNN, יכול לשנות איזה kernel נבחר - ולכן גם את הביצועים, לפעמים לרעה.
  • מצבי heuristic. ה-frontend מאפשר לבקש מצב הוריסטי מהיר (heur_mode.A) לעומת מצב שמתלבט יותר, ואף למדוד מספר תוכניות ולבחור את המהירה בפועל - אבל את זהות ה-kernel הסופי עדיין לא אתם קובעים.

וזו בדיוק נקודת המפנה אל CUTLASS. כשאתם צריכים kernel ספציפי שאתם כן רוצים לשלוט בו - צורה חריגה, fusion שהספרייה לא מכירה, או מיצוי אחרון של Tensor Cores על ארכיטקטורה חדשה - הקופסה השחורה ההוריסטית כבר לא מספיקה, ואתם יורדים לרמת הבנייה. על כך ב-6.3.

מרכז הכובד - CNNs על Ampere מול Transformers על Hopper

הנקודה הכי חשובה למקם נכון בשיעור הזה היא ש-cuDNN אינה "הפתרון של NVIDIA לכל רשת עמוקה" באופן אחיד. יש לה מרכז כובד היסטורי, ו-NVIDIA עצמה מכוונת מפתחים למקום אחר עבור העבודות החדשות ביותר.

cuDNN נולדה כדי להאיץ רשתות קונבולוציה (CNNs), וההתמחות ההיסטורית והחזקה ביותר שלה היא שם - במיוחד על ארכיטקטורת Ampere (דור ה-A100). כל עולם ה-ResNet, ה-U-Net והדומים להם יושב בנוחות על cuDNN, וה-fusion של conv-bias-activation שראינו הוא בדיוק התבנית שחוזרת שם אינסוף פעמים.

אבל בעולם ה-Transformers על הארכיטקטורות החדשות - Hopper (H100) ו-Blackwell (B200) - NVIDIA מדגישה דווקא את ספריית CUTLASS, ולא את cuDNN, ככלי המועדף לכתיבת ה-kernels הקריטיים (במיוחד ה-GEMM וה-attention המותכים שממצים את ה-Tensor Cores של הדור החדש). המשמעות המעשית:

ציר cuDNN CUTLASS (שיעור 6.3)
סוג primitives מוכנים לקריאה toolkit לבניית kernels
בחירת kernel הוריסטיקה סגורה, לא אתם אתם, בקוד
מרכז כובד CNNs על Ampere Transformers על Hopper/Blackwell
קוד frontend פתוח, backend סגור קוד פתוח מלא, header-only C++

אין כאן "ספרייה טובה מול רעה" - יש כלים לשכבות שונות של הצורך. אם אתם בונים inference של CNN קלאסי, cuDNN דרך ה-framework תיתן לכם ביצועים מצוינים כמעט בחינם. אם אתם בחזית של kernel ל-Transformer על Hopper - fused attention בצורה לא-סטנדרטית, למשל - תמצאו את עצמכם ב-CUTLASS. השיעור הבא, 6.3, נכנס בדיוק לפער הזה: מדוע ומתי יורדים מספרייה סגורה שקוראים לה, ל-toolkit פתוח שבונים איתו.

איך frameworks יושבים מעל cuDNN

בפועל, כמעט אף אחד לא קורא ל-cuDNN ישירות. הצרכן הכמעט-בלעדי שלה הוא שכבת ה-frameworks - PyTorch, TensorFlow, JAX - שמשתמשת בה כ-backend למימוש פעולות הרשת. כשאתם כותבים ב-PyTorch שורה תמימה כמו:

import torch
import torch.nn.functional as F

x = torch.randn(32, 256, 56, 56, device="cuda", dtype=torch.float16)
w = torch.randn(512, 256, 3, 3, device="cuda", dtype=torch.float16)

y = F.conv2d(x, w, padding=1)   # one line in Python
y = F.relu(y)

מה שקורה מתחת הוא בדיוק המחסנית שתיארנו: ה-conv2d מתורגם על ידי PyTorch לקריאה אל cuDNN, שבונה גרף פעולות, בוחרת מימוש בהוריסטיקה, ואולי אף ממזגת אותו עם ה-relu שאחריו ל-kernel יחיד. חלוקת העבודה בין האחיות ניכרת גם כאן: אם השורה הייתה torch.matmul על שני tensors דו-ממדיים, PyTorch היה נוטה לנתב אותה אל cuBLAS; משום שזו קונבולוציה, היא הולכת אל cuDNN.

שווה להזכיר גם את מנגנוני האוטו-כוונון של ה-frameworks. ב-PyTorch, הדגל torch.backends.cudnn.benchmark = True אומר ל-cuDNN למדוד בפועל כמה מימושים בהרצה הראשונה של כל צורת קלט חדשה, ולזכור את המהיר. זה משפר throughput כשצורות הקלט קבועות, אבל דוחף עלות למדידה בהרצה הראשונה - ואם הצורות משתנות כל הזמן, הוא דווקא מזיק. זו דוגמה יפה לכך שההוריסטיקה הסגורה של cuDNN חשופה כלפי מעלה דרך כמה כפתורים בודדים, אבל את הבחירה הסופית של ה-kernel עדיין אתם לא שולטים בה.

מי שלומד ללא GPU של NVIDIA יכול לתרגל את כל זה ב-Google Colab (כרטיס T4 חינמי) או על ענן שכור (Modal, Lambda, מכונות ענן) - בכל הסביבות האלה PyTorch כבר מגיע עם cuDNN מקומפל, וכל conv2d שתריצו כבר עובר דרכה.

סיכום

  • cuDNN היא ספריית ה-primitives של NVIDIA לרשתות עמוקות, אחות ל-cuBLAS: cuBLAS נותנת אלגברה לינארית כללית (dense GEMM), ו-cuDNN נותנת את ה-primitives המתמחים - convolution, attention, normalization, pooling ו-activations.
  • הקונבולוציה היא הליבה ההיסטורית של cuDNN, אבל היום היא מכסה גם scaled dot-product attention בסגנון FlashAttention, matmul ונרמולים, כך שהיא רלוונטית גם ל-Transformers ולא רק ל-CNNs.
  • ה-Graph API המודרני הוא דקלרטיבי: מתארים גרף שלם של פעולות דרך frontend פתוח ב-Python/C++, מעל backend C סגור, ורק אז cuDNN מנתחת את הגרף השלם ובונה תוכנית ביצוע.
  • ה-fusion ממזג רצף כמו Convolution+Bias+ReLU ל-kernel יחיד ששומר את תוצאות הביניים ב-shared memory במקום להוריד אותן ל-global memory, וכך חוסך רוחב פס - רווח שהוא כולו על תעבורת זיכרון, לא על כמות החישוב.
  • ה-fusion מתאפשר רק בזכות הראייה הגלובלית של הגרף: פעולה בודדת לא יודעת מה בא אחריה, ולכן חייבת לכתוב את תוצאתה ל-global; רק כש-cuDNN רואה את כל הרצף היא יכולה להשאיר את הביניים על ה-chip.
  • לכל primitive יש מספר מימושים, ו-cuDNN בוחרת ביניהם בהוריסטיקה סגורה ("unknown") לפי ארכיטקטורת ה-SM, ה-data type וגודל הקלט - אותה פילוסופיה אטומה של cuBLAS, על המחיר של ביצועים לא-צפויים ותלות בגרסה ובכרטיס.
  • מרכז הכובד של cuDNN הוא CNNs על Ampere; עבור Transformers על Hopper ו-Blackwell מפנה NVIDIA את המפתחים אל CUTLASS, וזו נקודת המעבר לשיעור 6.3.
  • בפועל כמעט אף אחד לא קורא ל-cuDNN ישירות: PyTorch, TensorFlow ו-JAX משתמשים בה כ-backend, ומנתבים קונבולוציות אליה בעוד כפלי מטריצות צפופים הולכים אל cuBLAS.