diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index 36a7e532aa..93ac53805f 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -28,6 +28,7 @@ jobs: - name: 'Build docs' env: NVTE_DOCS_DISABLE_PUBLIC_FEATURES: '1' + NVTE_DOCS_SWITCHER_JSON_URL: 'versions1.json' run: | # SPHINXOPTS="-W" errors out on warnings doxygen docs/Doxyfile cd docs diff --git a/docs/_static/css/diagram-colors.css b/docs/_static/css/diagram-colors.css index 96a2a8a6dc..f5dc7da4dd 100644 --- a/docs/_static/css/diagram-colors.css +++ b/docs/_static/css/diagram-colors.css @@ -70,6 +70,27 @@ font-size: 15px; font-weight: 600; text-anchor: middle; + fill: #212121; +} + +/* Accent text colors (used together with .text/.small-text/.section-title) */ +.accent-blue { fill: #1565c0; } +.accent-red { fill: #c62828; } +.accent-green { fill: #388e3c; } +.accent-orange { fill: #f57c00; } + +/* Neutral container panels (master_weights_approaches.svg) */ +.panel { + fill: #f5f5f5; + stroke: #9e9e9e; + stroke-width: 1.5; +} + +/* Range bracket lines (fp8_scaling_concept.svg) */ +.range-bracket { + stroke: #666; + stroke-width: 1.5; + fill: none; } /* Arrows */ @@ -132,3 +153,129 @@ stroke-width: 2.5; } +/* Dark mode overrides: darken fills and lighten text so diagrams stay + legible against the theme's dark background. */ +html[data-theme="dark"] .hp { fill: #3a2f5c; stroke: #b39ddb; } +html[data-theme="dark"] .fp8 { fill: #4d3c00; stroke: #ffca28; } +html[data-theme="dark"] .gemm { fill: #663c00; stroke: #ffb74d; } +html[data-theme="dark"] .quantize { fill: #1b3d1e; stroke: #81c784; } +html[data-theme="dark"] .amax { fill: #0d3c54; stroke: #4fc3f7; } +html[data-theme="dark"] .allgather { fill: #0d3c54; stroke: #4fc3f7; } +html[data-theme="dark"] .box-blue { fill: #10375c; stroke: #64b5f6; } +html[data-theme="dark"] .box-orange { fill: #5c3a10; stroke: #ffb74d; } +html[data-theme="dark"] .box-green { fill: #1e4620; stroke: #81c784; } +html[data-theme="dark"] .layernorm { fill: #10405c; stroke: #4fc3f7; } +html[data-theme="dark"] .fused { fill: #103c38; stroke: #4db6ac; } +html[data-theme="dark"] .computation { fill: #2b2b2b; stroke: #9e9e9e; } +html[data-theme="dark"] .fp32 { fill: #5c1f1f; stroke: #ef9a9a; } + +html[data-theme="dark"] .text, +html[data-theme="dark"] .title { + fill: #f5f5f5; +} + +html[data-theme="dark"] .small-text { + fill: #bdbdbd; +} + +html[data-theme="dark"] .label { + fill: #e0e0e0; +} + +html[data-theme="dark"] .arrow { + stroke: #bdbdbd; +} + +/* Text sitting directly on the page background (not inside boxes) */ +html[data-theme="dark"] .section-title, +html[data-theme="dark"] .section-label, +html[data-theme="dark"] .column-title, +html[data-theme="dark"] .dots-text, +html[data-theme="dark"] .dots-text-small { + fill: #e0e0e0; +} + +/* Lighter accents; must come after the .text/.small-text overrides */ +html[data-theme="dark"] .phase-forward, +html[data-theme="dark"] .accent-blue { fill: #64b5f6; } +html[data-theme="dark"] .phase-backward, +html[data-theme="dark"] .accent-red { fill: #ef9a9a; } +html[data-theme="dark"] .accent-green { fill: #81c784; } +html[data-theme="dark"] .accent-orange { fill: #ffb74d; } + +/* Arrowhead markers use fill attributes, which any CSS rule overrides */ +html[data-theme="dark"] marker polygon { + fill: #bdbdbd; +} + +/* Container panels and bracket/axis strokes on the page background */ +html[data-theme="dark"] .panel { fill: #1f1f1f; stroke: #757575; } +html[data-theme="dark"] .range-bracket { stroke: #bdbdbd; } +html[data-theme="dark"] .arrow-dashed { stroke: #bdbdbd; } +html[data-theme="dark"] .axis-line, +html[data-theme="dark"] .axis, +html[data-theme="dark"] .tick, +html[data-theme="dark"] .bwd-bracket { stroke: #bdbdbd; } + +/* Text on the page background in bit-layout / axis diagrams */ +html[data-theme="dark"] .format-label, +html[data-theme="dark"] .value-text, +html[data-theme="dark"] .tick-label { fill: #f5f5f5; } +html[data-theme="dark"] .header-text, +html[data-theme="dark"] .range-label, +html[data-theme="dark"] .sub-label { fill: #bdbdbd; } +html[data-theme="dark"] .param-label, +html[data-theme="dark"] .bracket-label { fill: #e0e0e0; } + +/* Flow-diagram box classes (mxfp8/blockwise swizzle-flow diagrams) */ +html[data-theme="dark"] .input-box, +html[data-theme="dark"] .quantize-box { fill: #3a2f5c; stroke: #b39ddb; } +html[data-theme="dark"] .quantize-fused-box { fill: #4a3a70; stroke: #b39ddb; } +html[data-theme="dark"] .mxfp8-box, +html[data-theme="dark"] .blockwise-box { fill: #10375c; stroke: #64b5f6; } +html[data-theme="dark"] .fp8-tile { fill: #16496f; stroke: #64b5f6; } +html[data-theme="dark"] .scale-tile, +html[data-theme="dark"] .gemm-box { fill: #1e4620; stroke: #81c784; } +html[data-theme="dark"] .scale-swizzled, +html[data-theme="dark"] .swizzle-box { fill: #5c3a10; stroke: #ffb74d; } +html[data-theme="dark"] .comm-box { fill: #4d3c00; stroke: #ffca28; } + +/* rht.svg and nvfp4_all_gather.svg boxes */ +html[data-theme="dark"] .input { fill: #10375c; stroke: #64b5f6; } +html[data-theme="dark"] .grad { fill: #5c1f38; stroke: #f48fb1; } +html[data-theme="dark"] .rht { fill: #5c3a10; stroke: #ffb74d; } +html[data-theme="dark"] .output { fill: #1e4620; stroke: #81c784; } +html[data-theme="dark"] .nvfp4 { fill: #10375c; stroke: #64b5f6; } + +/* White backing rects behind "..." glyphs: match the page background */ +html[data-theme="dark"] .diagram-bg { fill: #121212; } +html[data-theme="dark"] .tensor-outline { stroke: #9e9e9e; } + +/* Inter-frame connector segments drawn on the page background */ +html[data-theme="dark"] .connector-ext { stroke: #bdbdbd; } + +/* Ellipsis glyphs that sit on cells that stay light in dark mode */ +html[data-theme="dark"] .on-light-bg { fill: #212121; } + +/* cpu_offloading diagrams (classes unique to those four SVGs) */ +html[data-theme="dark"] .layer-box { fill: #2b2b2b; stroke: #9e9e9e; } +html[data-theme="dark"] .layer-text { fill: #e0e0e0; } +html[data-theme="dark"] .tensor-box { fill: #10375c; stroke: #64b5f6; } +html[data-theme="dark"] .tensor-text { fill: #90caf9; } +html[data-theme="dark"] .fwd-block { fill: #2e4a63; stroke: #7ba7cc; } +html[data-theme="dark"] .bwd-block { fill: #453a52; stroke: #b39ddb; } +html[data-theme="dark"] .offload-block { fill: #5c3a10; stroke: #d9a066; } +html[data-theme="dark"] .reload-block { fill: #2a4a26; stroke: #81c784; } +html[data-theme="dark"] .model-block { fill: #333333; stroke: #9e9e9e; } +html[data-theme="dark"] .wait-block { fill: #4a1d1d; stroke: #ef9a9a; } +html[data-theme="dark"] .wait-text { fill: #ef9a9a; } +html[data-theme="dark"] .block-text { fill: #e0e0e0; } +html[data-theme="dark"] .cpu-box { fill: #2e4a63; stroke: #7ba7cc; } +html[data-theme="dark"] .gpu-box { fill: #2a4a26; stroke: #81c784; } +html[data-theme="dark"] .frame-box { fill: #1f1f1f; stroke: #616161; } +html[data-theme="dark"] .memory-box { fill: #333c47; stroke: #9e9e9e; } +html[data-theme="dark"] .hbm-box { fill: #1e4620; stroke: #68d391; } +html[data-theme="dark"] .subtitle, +html[data-theme="dark"] .memory-label { fill: #e0e0e0; } +html[data-theme="dark"] .connector { stroke: #bdbdbd; } + diff --git a/docs/_static/css/output-style.css b/docs/_static/css/output-style.css index 864d8587a3..d32082e7df 100644 --- a/docs/_static/css/output-style.css +++ b/docs/_static/css/output-style.css @@ -1,10 +1,12 @@ /* Custom styling for program output blocks */ .program-output { - background-color: #f8f9fa; - padding: 0; /* No padding at all */ - margin: 0; /* No margins at all */ - border-radius: 0; /* No rounded corners */ + background-color: #f5f6f8; + border: 1px solid #d7dbe0; + border-top: none; + border-radius: 0 0 4px 4px; + padding: 8px 14px; + margin: 0 0 1rem 0; font-family: 'Courier New', monospace; font-size: 14px; line-height: 1.5; @@ -12,6 +14,18 @@ max-width: 100%; } +/* "Output:" bar attached between a code block and its program output */ +.output-header { + background: #eceef1; + border: 1px solid #d7dbe0; + border-left: 3px solid #9ca3af; + padding: 4px 12px; + font-size: 12px; + color: #4b5563; + margin-top: -16px; + margin-bottom: 0; +} + .program-output pre { margin: 0; padding: 0; @@ -27,6 +41,14 @@ width: 100%; } +.program-output div[class*="highlight"] { + margin: 0; +} + +.program-output pre { + padding: 2px 4px; +} + /* Alternative lighter style */ .output-block { background-color: #fafbfc; @@ -58,3 +80,62 @@ background: transparent !important; } +/* Header bar attached to the top of a code block ("Needs to be run on...") */ +.code-block-header { + background: #f0f4f8; + border-left: 3px solid #5c7cfa; + padding: 6px 12px; + font-size: 13px; + color: #495057; + margin-bottom: 0; + border-radius: 4px 4px 0 0; +} + +.code-block-header + div[class*="highlight"] { + margin-top: 0; +} + +/* Dark mode overrides */ +html[data-theme="dark"] .code-block-header { + background: #253141; + border-left-color: #748ffc; + color: #ced4da; +} + +html[data-theme="dark"] .program-output { + background-color: #232323; + border-color: #444950; +} + +html[data-theme="dark"] .output-header { + background: #2f3136; + border-color: #444950; + border-left-color: #757575; + color: #bdbdbd; +} + +html[data-theme="dark"] .program-output pre { + color: #e0e0e0; +} + +html[data-theme="dark"] .output-block { + background-color: #2b2b2b; + border-color: #444950; + color: #e0e0e0; +} + +html[data-theme="dark"] .console-output { + box-shadow: 0 2px 4px rgba(0, 0, 0, 0.4); +} + +/* pydata forces a light bg on rendered_html outputs in dark mode; + ours are pygments code blocks, so darken (same selector, loaded later) */ +html[data-theme="dark"] .bd-content .nboutput .output_area.rendered_html:not(:has(table.dataframe)) { + background-color: #1f1f1f; + color: #e0e0e0; +} + +html[data-theme="dark"] .output_area.stderr { + background: #4a1d1d; +} + diff --git a/docs/_static/css/sphinx_tabs.css b/docs/_static/css/sphinx_tabs.css index ecb2e2e00f..cfdde2a4d2 100644 --- a/docs/_static/css/sphinx_tabs.css +++ b/docs/_static/css/sphinx_tabs.css @@ -32,3 +32,25 @@ background-color: #f9f9f9; } +/* Dark mode overrides */ +html[data-theme="dark"] .sphinx-tabs-tab { + background-color: #2b2b2b; + border-color: #444950; + color: #e0e0e0; +} + +html[data-theme="dark"] .sphinx-tabs-tab:hover { + background-color: #3a3f44; +} + +html[data-theme="dark"] .sphinx-tabs-tab[aria-selected="true"] { + background-color: #76b900; + color: #1a1a1a; + border-color: #76b900; +} + +html[data-theme="dark"] .sphinx-tabs-panel { + background-color: #1f1f1f; + border-color: #444950; +} + diff --git a/docs/_static/css/svg-responsive.css b/docs/_static/css/svg-responsive.css index 4f0ebfa408..dbec1c5652 100644 --- a/docs/_static/css/svg-responsive.css +++ b/docs/_static/css/svg-responsive.css @@ -1,6 +1,13 @@ /* Responsive styling for project-authored SVG images */ -/* Keep embedded SVG content responsive without affecting theme UI icons. */ +/* Keep embedded SVG content responsive without affecting theme UI icons. + nvidia_sphinx_theme/pydata-sphinx-theme wraps page content in +
, not the RTD theme's
, + and `.. raw:: html :file:` doesn't add a `.raw-html` wrapper -- so these + selectors must target the article directly or SVGs with a fixed + width/viewBox will overflow the (narrower, "On this page" sidebar-adjacent) + content column. */ +article.bd-article svg, .document object[type="image/svg+xml"], .document .raw-html svg, .svg-figure svg { @@ -13,7 +20,16 @@ width: 100%; } +/* CSS cannot recolor -embedded SVGs: invert lightness instead, + hue-rotate(180deg) flips hue back so colors keep their tint */ +html[data-theme="dark"] article.bd-article img[src$=".svg"] { + filter: invert(0.88) hue-rotate(180deg); + /* undo pydata's white img background (inverts into a gray card) */ + background: transparent; +} + /* Ensure container doesn't overflow */ +article.bd-article, .document .raw-html { max-width: 100%; overflow-x: auto; diff --git a/docs/examples/jax/attention_context_parallel.rst b/docs/examples/jax/attention_context_parallel.rst index 5c2e2dd7ec..f701fe22af 100644 --- a/docs/examples/jax/attention_context_parallel.rst +++ b/docs/examples/jax/attention_context_parallel.rst @@ -125,7 +125,7 @@ while keeping those sharding controls visible. .. raw:: html -
+
Single-GPU output:
@@ -138,7 +138,7 @@ while keeping those sharding controls visible. .. raw:: html -
+
Ring output:
@@ -151,7 +151,7 @@ while keeping those sharding controls visible. .. raw:: html -
+
AllGather output:
diff --git a/docs/examples/jax/attention_single_gpu.rst b/docs/examples/jax/attention_single_gpu.rst index f675ebbe05..2c9f5dd8c8 100644 --- a/docs/examples/jax/attention_single_gpu.rst +++ b/docs/examples/jax/attention_single_gpu.rst @@ -69,7 +69,7 @@ implementations. .. raw:: html -
+
Output:
@@ -104,7 +104,7 @@ head dimension separately. .. raw:: html -
+
Output:
diff --git a/docs/examples/jax/dense.rst b/docs/examples/jax/dense.rst index 2087d49c7f..2bde67cd4c 100644 --- a/docs/examples/jax/dense.rst +++ b/docs/examples/jax/dense.rst @@ -64,7 +64,7 @@ same input for both models. .. raw:: html -
+
Output:
@@ -134,7 +134,7 @@ Two pieces wire this up: .. raw:: html -
+
Output:
diff --git a/docs/examples/te_gemma/media/calibration.svg b/docs/examples/te_gemma/media/calibration.svg index 16e1a43141..45629f1c9a 100644 --- a/docs/examples/te_gemma/media/calibration.svg +++ b/docs/examples/te_gemma/media/calibration.svg @@ -47,13 +47,7 @@ id="g62" clip-path="url(#clipPath62)" transform="translate(0,-153.77762)"> - + - + - + - + - + \ No newline at end of file + \ No newline at end of file diff --git a/docs/examples/te_llama/media/llama_zoom.svg b/docs/examples/te_llama/media/llama_zoom.svg index 6134ecfe1c..ce186a6b82 100644 --- a/docs/examples/te_llama/media/llama_zoom.svg +++ b/docs/examples/te_llama/media/llama_zoom.svg @@ -1 +1 @@ - \ No newline at end of file + \ No newline at end of file diff --git a/docs/examples/te_llama/media/llamadecoderlayer.svg b/docs/examples/te_llama/media/llamadecoderlayer.svg index 189369917d..606be546eb 100644 --- a/docs/examples/te_llama/media/llamadecoderlayer.svg +++ b/docs/examples/te_llama/media/llamadecoderlayer.svg @@ -1 +1 @@ - \ No newline at end of file + \ No newline at end of file diff --git a/docs/examples/te_llama/media/model_change.svg b/docs/examples/te_llama/media/model_change.svg index 6f0bed1927..2125bfe3a6 100644 --- a/docs/examples/te_llama/media/model_change.svg +++ b/docs/examples/te_llama/media/model_change.svg @@ -1 +1 @@ - \ No newline at end of file + \ No newline at end of file diff --git a/docs/examples/te_llama/media/swiglu.svg b/docs/examples/te_llama/media/swiglu.svg index 75b0a277a6..43dbbf6d48 100644 --- a/docs/examples/te_llama/media/swiglu.svg +++ b/docs/examples/te_llama/media/swiglu.svg @@ -1 +1 @@ - \ No newline at end of file + \ No newline at end of file diff --git a/docs/examples/te_llama/media/swiglu_te.svg b/docs/examples/te_llama/media/swiglu_te.svg index 5a846f2a0b..2808f9e4aa 100644 --- a/docs/examples/te_llama/media/swiglu_te.svg +++ b/docs/examples/te_llama/media/swiglu_te.svg @@ -1 +1 @@ - \ No newline at end of file + \ No newline at end of file diff --git a/docs/examples/te_llama/media/tellamadecoderlayer.svg b/docs/examples/te_llama/media/tellamadecoderlayer.svg index f93f49b720..f7889f48f3 100644 --- a/docs/examples/te_llama/media/tellamadecoderlayer.svg +++ b/docs/examples/te_llama/media/tellamadecoderlayer.svg @@ -1 +1 @@ - \ No newline at end of file + \ No newline at end of file diff --git a/docs/examples/te_llama/media/transformer_vs_llama.svg b/docs/examples/te_llama/media/transformer_vs_llama.svg index a872d6edec..c0cfca196b 100644 --- a/docs/examples/te_llama/media/transformer_vs_llama.svg +++ b/docs/examples/te_llama/media/transformer_vs_llama.svg @@ -1 +1 @@ - \ No newline at end of file + \ No newline at end of file diff --git a/docs/examples/te_llama/media/weight_swap.svg b/docs/examples/te_llama/media/weight_swap.svg index b2ff3ddf23..d6446800d3 100644 --- a/docs/examples/te_llama/media/weight_swap.svg +++ b/docs/examples/te_llama/media/weight_swap.svg @@ -1 +1 @@ - \ No newline at end of file + \ No newline at end of file diff --git a/docs/features/low_precision_training/fp8_blockwise_scaling/img/combined_scaling.svg b/docs/features/low_precision_training/fp8_blockwise_scaling/img/combined_scaling.svg index dbf6999aef..81676cad45 100644 --- a/docs/features/low_precision_training/fp8_blockwise_scaling/img/combined_scaling.svg +++ b/docs/features/low_precision_training/fp8_blockwise_scaling/img/combined_scaling.svg @@ -32,8 +32,8 @@ - - + + … @@ -47,7 +47,7 @@ … - + @@ -64,8 +64,8 @@ - - + + @@ -135,7 +135,7 @@ - + @@ -144,8 +144,8 @@ - - + + @@ -196,7 +196,7 @@ - + … @@ -221,8 +221,8 @@ - - + + @@ -276,7 +276,7 @@ - + @@ -285,8 +285,8 @@ - - + + @@ -335,7 +335,7 @@ - + Scaling factors (1 per 2D block) diff --git a/docs/features/low_precision_training/fp8_blockwise_scaling/img/transpose_handling.svg b/docs/features/low_precision_training/fp8_blockwise_scaling/img/transpose_handling.svg index e9a3b7b7d1..fa70565b45 100644 --- a/docs/features/low_precision_training/fp8_blockwise_scaling/img/transpose_handling.svg +++ b/docs/features/low_precision_training/fp8_blockwise_scaling/img/transpose_handling.svg @@ -24,8 +24,8 @@ - - + + @@ -96,7 +96,7 @@ - + @@ -107,8 +107,8 @@ - - + + @@ -187,7 +187,7 @@ - + @@ -207,8 +207,8 @@ - - + + @@ -262,7 +262,7 @@ - + @@ -279,8 +279,8 @@ - - + + @@ -340,7 +340,7 @@ - + diff --git a/docs/features/low_precision_training/fp8_current_scaling/fp8_current_scaling.rst b/docs/features/low_precision_training/fp8_current_scaling/fp8_current_scaling.rst index a4830a3fd5..cac3792194 100644 --- a/docs/features/low_precision_training/fp8_current_scaling/fp8_current_scaling.rst +++ b/docs/features/low_precision_training/fp8_current_scaling/fp8_current_scaling.rst @@ -139,7 +139,7 @@ Here's how to use FP8 Current Scaling recipe in PyTorch and JAX: .. raw:: html -
+
Requires SM89 (Ada) or later
@@ -152,7 +152,7 @@ Here's how to use FP8 Current Scaling recipe in PyTorch and JAX: .. raw:: html -
+
Requires SM89 (Ada) or later
diff --git a/docs/features/low_precision_training/fp8_current_scaling/img/fp8_scaling_concept.svg b/docs/features/low_precision_training/fp8_current_scaling/img/fp8_scaling_concept.svg index 9442b4e4aa..c7ad57f108 100644 --- a/docs/features/low_precision_training/fp8_current_scaling/img/fp8_scaling_concept.svg +++ b/docs/features/low_precision_training/fp8_current_scaling/img/fp8_scaling_concept.svg @@ -16,7 +16,7 @@ - + 0 @@ -34,12 +34,12 @@ - amax + amax - - - + + + Original range @@ -52,18 +52,18 @@ - + 0 - FP8 range + FP8 range - - FP8 range max + - FP8 range max @@ -83,14 +83,14 @@ - + 0 - FP8 range + FP8 range diff --git a/docs/features/low_precision_training/fp8_delayed_scaling/fp8_delayed_scaling.rst b/docs/features/low_precision_training/fp8_delayed_scaling/fp8_delayed_scaling.rst index 9d05305eda..d39787f6f5 100644 --- a/docs/features/low_precision_training/fp8_delayed_scaling/fp8_delayed_scaling.rst +++ b/docs/features/low_precision_training/fp8_delayed_scaling/fp8_delayed_scaling.rst @@ -89,7 +89,7 @@ Here's how to use FP8 Delayed Scaling in PyTorch and JAX: .. raw:: html -
+
Requires SM89 (Ada) or later
@@ -102,7 +102,7 @@ Here's how to use FP8 Delayed Scaling in PyTorch and JAX: .. raw:: html -
+
Requires SM89 (Ada) or later
diff --git a/docs/features/low_precision_training/introduction/img/fp8_linear_flow.svg b/docs/features/low_precision_training/introduction/img/fp8_linear_flow.svg index e1861ebc1c..48f020e137 100644 --- a/docs/features/low_precision_training/introduction/img/fp8_linear_flow.svg +++ b/docs/features/low_precision_training/introduction/img/fp8_linear_flow.svg @@ -13,7 +13,7 @@ FP8 Linear Layer – Forward and Backward Pass - Forward Pass + Forward Pass @@ -24,23 +24,23 @@ Input - + Quantize - + - + Input - + N @@ -48,16 +48,16 @@ Weight - + Quantize - + - + @@ -68,7 +68,7 @@ WeightT - + T @@ -77,7 +77,7 @@ (TN) - + @@ -87,7 +87,7 @@ - Backward Pass + Backward Pass @@ -98,16 +98,16 @@ Output grad. - + Quantize - + - + @@ -141,22 +141,22 @@ - + N - + T - + - + N - + T - + diff --git a/docs/features/low_precision_training/introduction/img/master_weights_approaches.svg b/docs/features/low_precision_training/introduction/img/master_weights_approaches.svg index b231fefd90..e1ff03e27e 100644 --- a/docs/features/low_precision_training/introduction/img/master_weights_approaches.svg +++ b/docs/features/low_precision_training/introduction/img/master_weights_approaches.svg @@ -24,7 +24,7 @@ (no master weights) - + Model Weights (BF16/FP16) @@ -40,7 +40,7 @@ - + Optimizer State (FP32) @@ -49,7 +49,7 @@ Master Weights in Model - + Model Weights (FP32) @@ -67,7 +67,7 @@ - + Optimizer State (FP32) @@ -83,7 +83,7 @@ - + Model Weights (BF16/FP16) @@ -99,7 +99,7 @@ - + Optimizer State (FP32) diff --git a/docs/features/low_precision_training/introduction/introduction.rst b/docs/features/low_precision_training/introduction/introduction.rst index 760a63b0b1..fba7796ece 100644 --- a/docs/features/low_precision_training/introduction/introduction.rst +++ b/docs/features/low_precision_training/introduction/introduction.rst @@ -143,7 +143,7 @@ Let's now see how we can train in lower precisions in supported frameworks. .. raw:: html -
+
Needs to be run on SM89+ (Ada or newer)
@@ -158,7 +158,7 @@ Let's now see how we can train in lower precisions in supported frameworks. .. raw:: html -
+
Needs to be run on SM89+ (Ada or newer)
@@ -171,7 +171,7 @@ Let's now see how we can train in lower precisions in supported frameworks. .. raw:: html -
+
Needs to be run on SM89+ (Ada or newer)
@@ -191,7 +191,7 @@ Let's now see how we can train in lower precisions in supported frameworks. .. raw:: html -
+
Needs to be run on SM89+ (Ada or newer)
@@ -206,7 +206,7 @@ Let's now see how we can train in lower precisions in supported frameworks. .. raw:: html -
+
Needs to be run on SM89+ (Ada or newer)
@@ -219,7 +219,7 @@ Let's now see how we can train in lower precisions in supported frameworks. .. raw:: html -
+
Needs to be run on SM89+ (Ada or newer)
diff --git a/docs/features/low_precision_training/mxfp8/img/fp8_1d_scaling.svg b/docs/features/low_precision_training/mxfp8/img/fp8_1d_scaling.svg index 30f16d9a71..f1f37cdb57 100644 --- a/docs/features/low_precision_training/mxfp8/img/fp8_1d_scaling.svg +++ b/docs/features/low_precision_training/mxfp8/img/fp8_1d_scaling.svg @@ -29,8 +29,8 @@ - - + + @@ -100,7 +100,7 @@ - + @@ -109,8 +109,8 @@ - - + + @@ -161,7 +161,7 @@ - + … diff --git a/docs/features/low_precision_training/mxfp8/img/mxfp8_row_col.svg b/docs/features/low_precision_training/mxfp8/img/mxfp8_row_col.svg index 42ea0308bb..ed28fa456b 100644 --- a/docs/features/low_precision_training/mxfp8/img/mxfp8_row_col.svg +++ b/docs/features/low_precision_training/mxfp8/img/mxfp8_row_col.svg @@ -73,7 +73,7 @@ - + … … … @@ -141,7 +141,7 @@ - + … … … @@ -200,7 +200,7 @@ - + … … … @@ -258,7 +258,7 @@ - + … … … diff --git a/docs/features/low_precision_training/mxfp8/img/mxfp8_scale_linearize_and_swizzle.svg b/docs/features/low_precision_training/mxfp8/img/mxfp8_scale_linearize_and_swizzle.svg index 6e4ed44d56..ddf8cb2f5c 100644 --- a/docs/features/low_precision_training/mxfp8/img/mxfp8_scale_linearize_and_swizzle.svg +++ b/docs/features/low_precision_training/mxfp8/img/mxfp8_scale_linearize_and_swizzle.svg @@ -23,7 +23,7 @@ - + @@ -72,14 +72,14 @@ 128x4 - + - + - + @@ -89,7 +89,7 @@ 2 - … + … K + 1 @@ -102,7 +102,7 @@ 1x512 - + @@ -118,10 +118,10 @@ Q 250 345, 255 345 L 375 345 Q 385 345, 385 355" - fill="none" stroke="#444" stroke-width="2"/> + class="tensor-outline" fill="none" stroke="#444" stroke-width="2"/> - + @@ -152,11 +152,11 @@ 5 6 7 - ... + ... - + @@ -183,7 +183,7 @@ 33 65 97 - ... + ... diff --git a/docs/features/low_precision_training/mxfp8/img/mxfp8_tensor_scaling_layout.svg b/docs/features/low_precision_training/mxfp8/img/mxfp8_tensor_scaling_layout.svg index 3b81ff0a36..f5b24568dc 100644 --- a/docs/features/low_precision_training/mxfp8/img/mxfp8_tensor_scaling_layout.svg +++ b/docs/features/low_precision_training/mxfp8/img/mxfp8_tensor_scaling_layout.svg @@ -9,7 +9,7 @@ - + @@ -37,7 +37,7 @@ - + diff --git a/docs/features/low_precision_training/mxfp8/mxfp8.rst b/docs/features/low_precision_training/mxfp8/mxfp8.rst index f8f8f48b0d..1fbcc43af9 100644 --- a/docs/features/low_precision_training/mxfp8/mxfp8.rst +++ b/docs/features/low_precision_training/mxfp8/mxfp8.rst @@ -121,7 +121,7 @@ Here's how to use MXFP8 recipe in PyTorch and JAX: .. raw:: html -
+
Requires SM100 (Blackwell) or later
@@ -134,7 +134,7 @@ Here's how to use MXFP8 recipe in PyTorch and JAX: .. raw:: html -
+
Requires SM100 (Blackwell) or later
diff --git a/docs/features/low_precision_training/nvfp4/img/nvfp4_hierarchical_scaling.svg b/docs/features/low_precision_training/nvfp4/img/nvfp4_hierarchical_scaling.svg index 05e67b7889..dc9a3a96a8 100644 --- a/docs/features/low_precision_training/nvfp4/img/nvfp4_hierarchical_scaling.svg +++ b/docs/features/low_precision_training/nvfp4/img/nvfp4_hierarchical_scaling.svg @@ -30,8 +30,8 @@ - - + + @@ -101,7 +101,7 @@ - + @@ -110,8 +110,8 @@ - - + + @@ -162,7 +162,7 @@ - + … diff --git a/docs/features/low_precision_training/nvfp4/img/nvfp4_row_col.svg b/docs/features/low_precision_training/nvfp4/img/nvfp4_row_col.svg index 30363d6ce2..48c6a553df 100644 --- a/docs/features/low_precision_training/nvfp4/img/nvfp4_row_col.svg +++ b/docs/features/low_precision_training/nvfp4/img/nvfp4_row_col.svg @@ -56,7 +56,7 @@ ⋮ ⋮ - + @@ -100,7 +100,7 @@ ⋮ ⋮ - + @@ -153,7 +153,7 @@ ⋮ ⋮ - + @@ -197,7 +197,7 @@ ⋮ ⋮ - + diff --git a/docs/features/low_precision_training/nvfp4/nvfp4.rst b/docs/features/low_precision_training/nvfp4/nvfp4.rst index 0415963a71..900789b0d3 100644 --- a/docs/features/low_precision_training/nvfp4/nvfp4.rst +++ b/docs/features/low_precision_training/nvfp4/nvfp4.rst @@ -218,7 +218,7 @@ Here's how to use NVFP4 recipe in PyTorch and JAX. The examples show how to conf .. raw:: html -
+
Requires SM100 (Blackwell) or later
@@ -231,7 +231,7 @@ Here's how to use NVFP4 recipe in PyTorch and JAX. The examples show how to conf .. raw:: html -
+
Requires SM100 (Blackwell) or later
diff --git a/docs/features/low_precision_training/performance_considerations/img/gemm_access_pattern.svg b/docs/features/low_precision_training/performance_considerations/img/gemm_access_pattern.svg index fa720427e7..ad84fa315a 100644 --- a/docs/features/low_precision_training/performance_considerations/img/gemm_access_pattern.svg +++ b/docs/features/low_precision_training/performance_considerations/img/gemm_access_pattern.svg @@ -7,6 +7,8 @@ .small-text { font: 9px sans-serif; fill: #555; } .operator { font: bold 14px sans-serif; fill: #333; } .equals { font: bold 14px sans-serif; fill: #333; } + .accent-green { fill: #388e3c; } + .accent-orange { fill: #f57c00; } /* Matrix colors */ .matrix-cell { fill: #bbdefb; stroke: #1565c0; stroke-width: 1; } @@ -50,7 +52,7 @@ - rowwise + rowwise @@ -82,7 +84,7 @@ - columnwise + columnwise @@ -148,7 +150,7 @@ - rowwise + rowwise @@ -180,7 +182,7 @@ - rowwise + rowwise diff --git a/docs/features/low_precision_training/performance_considerations/img/sequence_parallel_quantization.svg b/docs/features/low_precision_training/performance_considerations/img/sequence_parallel_quantization.svg index 5b61ac2478..99c8b3f2bc 100644 --- a/docs/features/low_precision_training/performance_considerations/img/sequence_parallel_quantization.svg +++ b/docs/features/low_precision_training/performance_considerations/img/sequence_parallel_quantization.svg @@ -24,10 +24,10 @@ - All-Gather of Quantized Tensors (one scenario) + All-Gather of Quantized Tensors (one scenario) - Input Tensor quantized all-gather + Input Tensor quantized all-gather FWD: @@ -93,7 +93,7 @@ - Gradient Tensor quantized all-gather + Gradient Tensor quantized all-gather BWD: diff --git a/docs/features/low_precision_training/performance_considerations/img/transpose_fusion.svg b/docs/features/low_precision_training/performance_considerations/img/transpose_fusion.svg index 194b1237e1..a1b9dd9e86 100644 --- a/docs/features/low_precision_training/performance_considerations/img/transpose_fusion.svg +++ b/docs/features/low_precision_training/performance_considerations/img/transpose_fusion.svg @@ -18,7 +18,7 @@ - Option 1: Quantize both usages in forward + Option 1: Quantize both usages in forward FORWARD: @@ -58,7 +58,7 @@ - Option 2: Separate Quantizations (quantize when needed) + Option 2: Separate Quantizations (quantize when needed) FORWARD: @@ -113,7 +113,7 @@ - Option 3: Convert Rowwise to Columnwise in Backward (reuse saved tensor) + Option 3: Convert Rowwise to Columnwise in Backward (reuse saved tensor) FORWARD: diff --git a/docs/features/low_precision_training/performance_considerations/performance_considerations.rst b/docs/features/low_precision_training/performance_considerations/performance_considerations.rst index a495af56c1..2c21799dd6 100644 --- a/docs/features/low_precision_training/performance_considerations/performance_considerations.rst +++ b/docs/features/low_precision_training/performance_considerations/performance_considerations.rst @@ -191,7 +191,7 @@ and columnwise tensors require separate memory layouts. .. raw:: html -
+
Needs to be run on SM89 (Ada) or SM90 (Hopper)
@@ -202,7 +202,7 @@ and columnwise tensors require separate memory layouts. .. raw:: html -
+
Output:
@@ -222,7 +222,7 @@ and columnwise tensors require separate memory layouts. .. raw:: html -
+
Needs to be run on SM89 (Ada) or SM90 (Hopper)
@@ -233,7 +233,7 @@ and columnwise tensors require separate memory layouts. .. raw:: html -
+
Output:
@@ -253,7 +253,7 @@ and columnwise tensors require separate memory layouts. .. raw:: html -
+
Needs to be run on SM89 (Ada) or SM90 (Hopper)
@@ -264,7 +264,7 @@ and columnwise tensors require separate memory layouts. .. raw:: html -
+
Output:
@@ -290,7 +290,7 @@ and columnwise tensors require separate memory layouts. .. raw:: html -
+
Needs to be run on SM89 (Ada) or SM90 (Hopper)
@@ -301,7 +301,7 @@ and columnwise tensors require separate memory layouts. .. raw:: html -
+
Output:
@@ -320,7 +320,7 @@ and columnwise tensors require separate memory layouts. .. raw:: html -
+
Needs to be run on SM89 (Ada) or SM90 (Hopper)
@@ -331,7 +331,7 @@ and columnwise tensors require separate memory layouts. .. raw:: html -
+
Output:
@@ -351,7 +351,7 @@ and columnwise tensors require separate memory layouts. .. raw:: html -
+
Needs to be run on SM89 (Ada) or SM90 (Hopper)
@@ -362,7 +362,7 @@ and columnwise tensors require separate memory layouts. .. raw:: html -
+
Output:
@@ -407,7 +407,7 @@ Let's see how we can use fused layers in different frameworks. .. raw:: html -
+
Needs to be run on SM89+ (Ada, Hopper, Blackwell, or newer)
@@ -427,7 +427,7 @@ Let's see how we can use fused layers in different frameworks. .. raw:: html -
+
Needs to be run on SM89+ (Ada, Hopper, Blackwell, or newer)
diff --git a/docs/features/other_optimizations/cpu_offloading/img/pcie_vs_nvlink.svg b/docs/features/other_optimizations/cpu_offloading/img/pcie_vs_nvlink.svg index 0b8ec3912a..d6f1210863 100644 --- a/docs/features/other_optimizations/cpu_offloading/img/pcie_vs_nvlink.svg +++ b/docs/features/other_optimizations/cpu_offloading/img/pcie_vs_nvlink.svg @@ -6,7 +6,7 @@ .label { font: 12px sans-serif; fill: #333; text-anchor: middle; } .memory-label { font: 12px sans-serif; fill: #333; text-anchor: middle; } .small-text { font: 11px sans-serif; fill: #555; text-anchor: middle; } - .speed-text { font: bold 11px sans-serif; fill: #fff; text-anchor: middle; } + .speed-text { font: bold 11px sans-serif; fill: #1a1a1a; text-anchor: middle; } .cpu-box { fill: #9db4d0; stroke: #444; stroke-width: 2; } .gpu-box { fill: #a8d99c; stroke: #444; stroke-width: 2; } @@ -60,7 +60,7 @@ - + PCIe @@ -89,7 +89,7 @@ - + NVLink C2C @@ -108,7 +108,7 @@ - + NVLink C2C diff --git a/docs/getting_started/index.rst b/docs/getting_started/index.rst index 9e10f82c14..a7ed14c2f4 100644 --- a/docs/getting_started/index.rst +++ b/docs/getting_started/index.rst @@ -54,7 +54,7 @@ for later comparisons with Transformer Engine. .. raw:: html -

Structure of a GPT decoder layer

+

Structure of a GPT decoder layer

We construct the components as follows: @@ -107,7 +107,7 @@ Putting it all together: .. raw:: html -
+
Output:
@@ -143,7 +143,7 @@ Putting it all together: .. raw:: html -
+
Output:
@@ -195,7 +195,7 @@ This is the simplest way to start using Transformer Engine. .. raw:: html -
+
Output:
@@ -239,7 +239,7 @@ This is the simplest way to start using Transformer Engine. .. raw:: html -
+
Output:
@@ -280,7 +280,7 @@ delivering optimal performance without manual tuning. .. raw:: html -
+
Output:
@@ -311,7 +311,7 @@ delivering optimal performance without manual tuning. .. raw:: html -
+
Output:
@@ -361,7 +361,7 @@ This provides significant speedups on supported hardware (Hopper, Ada, Blackwell .. raw:: html -
+
Output:
@@ -402,7 +402,7 @@ This provides significant speedups on supported hardware (Hopper, Ada, Blackwell .. raw:: html -
+
Output:
@@ -444,7 +444,7 @@ Combined with TE Attention and FP8, this delivers peak performance. .. raw:: html -
+
Output:
@@ -476,7 +476,7 @@ Combined with TE Attention and FP8, this delivers peak performance. .. raw:: html -
+
Output:
@@ -507,7 +507,7 @@ module that includes all optimizations out of the box. .. raw:: html -
+
Output:
@@ -529,7 +529,7 @@ module that includes all optimizations out of the box. .. raw:: html -
+
Output:
diff --git a/docs/getting_started/transformer_layer.svg b/docs/getting_started/transformer_layer.svg index 28ba3dd386..5e9040e24c 100644 --- a/docs/getting_started/transformer_layer.svg +++ b/docs/getting_started/transformer_layer.svg @@ -20,38 +20,38 @@ - + LayerNorm - + QKV Projection - + Dot Product Attention - + Projection - + Dropout - + + @@ -61,19 +61,19 @@ - + LayerNorm - + MLP - + +