diff --git a/.gitignore b/.gitignore index a6d79ddb5b..6edaa04783 100644 --- a/.gitignore +++ b/.gitignore @@ -92,6 +92,9 @@ instance/ docs/_build docs/apidocs +# Fern autodoc (generated by fern docs md generate) +fern/product-docs/ + # PyBuilder target/ diff --git a/3rdparty/Megatron-LM b/3rdparty/Megatron-LM index 8f1c2f8ae5..0d8e0714cd 160000 --- a/3rdparty/Megatron-LM +++ b/3rdparty/Megatron-LM @@ -1 +1 @@ -Subproject commit 8f1c2f8ae53b4e3f32c0ae7f397d8b38a675eaa2 +Subproject commit 0d8e0714cd29c01e164fe6de9f532182bdffa942 diff --git a/fern/README.md b/fern/README.md new file mode 100644 index 0000000000..8959b34bcd --- /dev/null +++ b/fern/README.md @@ -0,0 +1,100 @@ +# Megatron Bridge Fern Documentation + +This folder contains the Fern Docs configuration for Megatron Bridge. + +## Installation + +```bash +npm install -g fern-api +# Or: npx fern-api --version +``` + +## Local Preview + +```bash +cd fern/ +fern docs dev +# Or from project root: fern docs dev --project ./fern +``` + +Docs available at `http://localhost:3000`. + +## Folder Structure + +``` +fern/ +├── docs.yml # Global config (title, colors, versions) +├── fern.config.json # Fern CLI config +├── versions/ +│ └── v0.2.0.yml # Navigation for v0.2.0 +├── v0.2.0/ +│ └── pages/ # MDX content for v0.2.0 +├── scripts/ # Migration and conversion scripts +└── assets/ # Favicon, images +``` + +## Migration Workflow + +To migrate or update docs from `docs/` to Fern: + +```bash +# 1. Copy docs to fern (run from repo root) +python3 fern/scripts/copy_docs_to_fern.py v0.2.0 + +# 2. Expand {include} directives (index, changelog) +python3 fern/scripts/expand_includes.py fern/v0.2.0/pages + +# 3. Convert MB-specific syntax (py:class, py:meth) +python3 fern/scripts/convert_mb_specific.py fern/v0.2.0/pages + +# 4. Convert MyST to Fern MDX +python3 fern/scripts/convert_myst_to_fern.py fern/v0.2.0/pages + +# 5. Add frontmatter +python3 fern/scripts/add_frontmatter.py fern/v0.2.0/pages + +# 6. Update internal links +python3 fern/scripts/update_links.py fern/v0.2.0/pages + +# 7. Remove duplicate H1s (when title matches frontmatter) +python3 fern/scripts/remove_duplicate_h1.py fern/v0.2.0/pages + +# 8. Validate +./fern/scripts/check_unconverted.sh fern/v0.2.0/pages +``` + +## MDX Components + +```mdx +Informational note +Helpful tip +Warning message +Info callout + + + Description + + + + ```python\ncode\n``` + + +Collapsible content +``` + +## API Reference + +API docs are built by Sphinx (autodoc2) and hosted at docs.nvidia.com. The "API Reference" link in the navbar points to `https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/`. + +## Deploying + +```bash +fern generate --docs +fern docs deploy +``` + +## Useful Links + +- [Fern Docs](https://buildwithfern.com/learn/docs) +- [MDX Components](https://buildwithfern.com/learn/docs/components) +- [Versioning Guide](https://buildwithfern.com/learn/docs/configuration/versions) diff --git a/fern/assets/NVIDIA_dark.svg b/fern/assets/NVIDIA_dark.svg new file mode 100644 index 0000000000..04850d9d6b --- /dev/null +++ b/fern/assets/NVIDIA_dark.svg @@ -0,0 +1,35 @@ + + + + + + + + + + + + + + + + + + + diff --git a/fern/assets/NVIDIA_light.svg b/fern/assets/NVIDIA_light.svg new file mode 100644 index 0000000000..9ee045c3ef --- /dev/null +++ b/fern/assets/NVIDIA_light.svg @@ -0,0 +1,34 @@ + + + + + + + + + + + + + + + + + + + diff --git a/fern/assets/NVIDIA_symbol.svg b/fern/assets/NVIDIA_symbol.svg new file mode 100644 index 0000000000..c0507afe00 --- /dev/null +++ b/fern/assets/NVIDIA_symbol.svg @@ -0,0 +1,22 @@ + + + + + + + + + + + + + + + diff --git a/fern/assets/README.md b/fern/assets/README.md new file mode 100644 index 0000000000..dc0d33c50d --- /dev/null +++ b/fern/assets/README.md @@ -0,0 +1,4 @@ +# Fern Assets + +Add `favicon.png` here for the docs site logo and favicon. +See NeMo Curator or DataDesigner fern/assets for reference. diff --git a/fern/assets/Repo-Mbridge.png b/fern/assets/Repo-Mbridge.png new file mode 100644 index 0000000000..5c5dace2ad Binary files /dev/null and b/fern/assets/Repo-Mbridge.png differ diff --git a/fern/assets/favicon.png b/fern/assets/favicon.png new file mode 100644 index 0000000000..5c5dace2ad Binary files /dev/null and b/fern/assets/favicon.png differ diff --git a/fern/assets/training/images/activation-recomputation-example-1.jpg b/fern/assets/training/images/activation-recomputation-example-1.jpg new file mode 100644 index 0000000000..78322e6107 Binary files /dev/null and b/fern/assets/training/images/activation-recomputation-example-1.jpg differ diff --git a/fern/assets/training/images/activation-recomputation-example-2.jpg b/fern/assets/training/images/activation-recomputation-example-2.jpg new file mode 100644 index 0000000000..1376fa42f0 Binary files /dev/null and b/fern/assets/training/images/activation-recomputation-example-2.jpg differ diff --git a/fern/assets/training/images/canonical_lora.png b/fern/assets/training/images/canonical_lora.png new file mode 100644 index 0000000000..69e8dacf09 Binary files /dev/null and b/fern/assets/training/images/canonical_lora.png differ diff --git a/fern/assets/training/images/ddp.gif b/fern/assets/training/images/ddp.gif new file mode 100644 index 0000000000..08ac652e0c Binary files /dev/null and b/fern/assets/training/images/ddp.gif differ diff --git a/fern/assets/training/images/ep.png b/fern/assets/training/images/ep.png new file mode 100644 index 0000000000..66e868906e Binary files /dev/null and b/fern/assets/training/images/ep.png differ diff --git a/fern/assets/training/images/mtp_loss.png b/fern/assets/training/images/mtp_loss.png new file mode 100644 index 0000000000..bcbae1aa3c Binary files /dev/null and b/fern/assets/training/images/mtp_loss.png differ diff --git a/fern/assets/training/images/mtp_loss_comparison.png b/fern/assets/training/images/mtp_loss_comparison.png new file mode 100644 index 0000000000..b02db9c74f Binary files /dev/null and b/fern/assets/training/images/mtp_loss_comparison.png differ diff --git a/fern/assets/training/images/performant_lora.png b/fern/assets/training/images/performant_lora.png new file mode 100644 index 0000000000..00c12df247 Binary files /dev/null and b/fern/assets/training/images/performant_lora.png differ diff --git a/fern/assets/training/images/pnom.gif b/fern/assets/training/images/pnom.gif new file mode 100644 index 0000000000..bddfbd433a Binary files /dev/null and b/fern/assets/training/images/pnom.gif differ diff --git a/fern/assets/training/images/pp.gif b/fern/assets/training/images/pp.gif new file mode 100644 index 0000000000..d14bac5774 Binary files /dev/null and b/fern/assets/training/images/pp.gif differ diff --git a/fern/assets/training/images/pp_comm_overlap.png b/fern/assets/training/images/pp_comm_overlap.png new file mode 100644 index 0000000000..efaaf8f727 Binary files /dev/null and b/fern/assets/training/images/pp_comm_overlap.png differ diff --git a/fern/assets/training/images/sp.png b/fern/assets/training/images/sp.png new file mode 100644 index 0000000000..aba01deffd Binary files /dev/null and b/fern/assets/training/images/sp.png differ diff --git a/fern/assets/training/images/tp1.png b/fern/assets/training/images/tp1.png new file mode 100644 index 0000000000..1903488806 Binary files /dev/null and b/fern/assets/training/images/tp1.png differ diff --git a/fern/assets/training/images/tp2.png b/fern/assets/training/images/tp2.png new file mode 100644 index 0000000000..33f25b117f Binary files /dev/null and b/fern/assets/training/images/tp2.png differ diff --git a/fern/assets/training/images/tp_comm_overlap.png b/fern/assets/training/images/tp_comm_overlap.png new file mode 100644 index 0000000000..4b44b20a34 Binary files /dev/null and b/fern/assets/training/images/tp_comm_overlap.png differ diff --git a/fern/components/CustomFooter.tsx b/fern/components/CustomFooter.tsx new file mode 100644 index 0000000000..fab392c407 --- /dev/null +++ b/fern/components/CustomFooter.tsx @@ -0,0 +1,91 @@ +/** + * Custom footer for NVIDIA docs (Fern native header/footer). + * Markup and class names match the original custom-app footer 1:1 so that + * fern/main.css (footer + Built with Fern styles) applies correctly: + * dark mode logo, responsive layout, and Built with Fern tooltip. + */ +export default function CustomFooter() { + const currentYear = new Date().getFullYear(); + const logoUrl = + "https://fern-image-hosting.s3.us-east-1.amazonaws.com/nvidia/NVIDIA_Logo_0.svg"; + + return ( + + ); +} diff --git a/fern/docs.yml b/fern/docs.yml new file mode 100644 index 0000000000..873d18392b --- /dev/null +++ b/fern/docs.yml @@ -0,0 +1,63 @@ +instances: + - url: https://megatron-bridge.docs.buildwithfern.com + +title: Megatron Bridge + +libraries: + megatron-bridge: + input: + git: https://github.com/NVIDIA-NeMo/Megatron-Bridge + subpath: src/megatron/bridge + output: + path: ./product-docs/megatron-bridge/Library-Reference + lang: python + +versions: + - display-name: v0.2.0 + path: versions/v0.2.0.yml + slug: v0.2.0 + +footer: ./components/CustomFooter.tsx + +layout: + searchbar-placement: header + page-width: 1376px + sidebar-width: 248px + content-width: 812px + tabs-placement: header + hide-feedback: true + +colors: + accentPrimary: + dark: "#76B900" + light: "#76B900" + background: + light: "#FFFFFF" + dark: "#000000" + +theme: + page-actions: toolbar + footer-nav: minimal + +logo: + dark: ./assets/NVIDIA_dark.svg + light: ./assets/NVIDIA_light.svg + height: 20 + href: / + right-text: Megatron Bridge + +favicon: ./assets/NVIDIA_symbol.svg + +css: + - ./main.css + +navbar-links: + - type: github + value: https://github.com/NVIDIA-NeMo/Megatron-Bridge + - type: secondary + text: API Reference + url: https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/ + +experimental: + mdx-components: + - ./components diff --git a/fern/fern.config.json b/fern/fern.config.json new file mode 100644 index 0000000000..006b3f4577 --- /dev/null +++ b/fern/fern.config.json @@ -0,0 +1,4 @@ +{ + "organization": "nvidia", + "version": "3.79.3" +} diff --git a/fern/main.css b/fern/main.css new file mode 100644 index 0000000000..87f5dbf7f9 --- /dev/null +++ b/fern/main.css @@ -0,0 +1,867 @@ +/*! + * SPDX-FileCopyrightText: Copyright (c) 2023-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: LicenseRef-NvidiaProprietary + * + * NVIDIA CORPORATION, its affiliates and licensors retain all intellectual + * property and proprietary rights in and to this material, related + * documentation and any modifications thereto. Any use, reproduction, + * disclosure or distribution of this material and related documentation + * without an express license agreement from NVIDIA CORPORATION or + * its affiliates is strictly prohibited. + */ + +/* Color themes for light and dark modes */ +:root { + /* Brand Colors */ + --nv-color-green: #74B900; + --nv-color-green-2: #004B31; + --nv-color-black: #000000; + --nv-color-white: #FFFFFF; + + /* Grey Scale - Light */ + --nv-light-grey-1: #f7f7f7; + --nv-light-grey-2: #EEEEEE; + --nv-light-grey-3: #DDDDDD; + --nv-light-grey-4: #CCCCCC; + --nv-light-grey-5: #999999; + + /* Grey Scale - Dark */ + --nv-dark-grey-1: #111111; + --nv-dark-grey-2: #1A1A1A; + --nv-dark-grey-3: #222222; + --nv-dark-grey-4: #333333; + --nv-dark-grey-5: #666666; + + /* Colors by Usage */ + --nv-color-text: #000000; + --nv-color-bg-default: #FFFFFF; + --nv-color-bg-alt: #f7f7f7; + --nv-color-success: #76B900; + --nv-color-error: #f44336; + + /* Theme-independent settings */ + --rounded: 999px; +} +main { + min-height: calc(100vh - 200px); + } +/* Typography - Headers */ +h1 { + font-size: 36px; + font-weight: 700; + line-height: 1.25em; /* 45px */ +} + +h2 { + font-size: 28px; + font-weight: 700; + line-height: 1.25em; /* 35px */ +} + +h3 { + font-size: 24px; + font-weight: 700; + line-height: 1.25em; /* 30px */ +} + +h4 { + font-size: 20px; + font-weight: 700; + line-height: 1.25em; /* 25px */ +} + +/* Typography - Paragraphs */ +.prose{ + color: var(--nv-dark-grey-2) !important; +} +.dark .prose{ + color: var(--nv-light-grey-2) !important; +} +p { + text-decoration-thickness: 3px; +} +.fern-mdx-link { + color: var(--tw-prose-body); + text-decoration-color: var(--accent); + font-weight: var(--font-weight-normal); +} + +/* Light theme (default) */ +html:not([data-theme]),html[data-theme=light] { + --pst-color-background: #fff; + --pst-color-on-background: #fff; + --pst-color-shadow: #ccc; + --pst-color-heading: #000; + --pst-color-text-base: #1a1a1a; + --pst-color-text-muted: #666; + --pst-color-surface: #f7f7f7; + --pst-color-on-surface: #333; + --pst-color-primary: var(--nv-color-green-2); + --pst-color-table-row-hover-bg: var(--nv-color-green); + --pst-color-link: var(--pst-color-text-base); + --pst-color-link-hover: var(--pst-color-text-base); + --pst-color-inline-code: var(--pst-color-primary); + --pst-color-inline-code-links: var(--pst-color-primary); + --pst-color-secondary: var(--pst-color-primary); + --pst-color-secondary-bg: var(--nv-color-green); + --pst-color-accent: var(--nv-color-green); +} + +/* Dark theme */ +html[data-theme=dark] { + --pst-color-background: #111; + --pst-color-on-background: #000; + --pst-color-shadow: #000; + --pst-color-heading: #fff; + --pst-color-text-base: #eee; + --pst-color-text-muted: #999; + --pst-color-surface: #1a1a1a; + --pst-color-on-surface: #ddd; + --pst-color-primary: var(--nv-color-green); + --pst-color-table-row-hover-bg: var(--nv-color-green-2); + --pst-color-link: var(--pst-color-text-base); + --pst-color-link-hover: var(--pst-color-text-base); + --pst-color-inline-code: var(--pst-color-primary); + --pst-color-inline-code-links: var(--pst-color-primary); + --pst-color-secondary: var(--pst-color-primary); + --pst-color-secondary-bg: var(--nv-color-green-2); + --pst-color-accent: var(--nv-color-green); +} + +/* Product and verion selector styling */ + +.fern-product-selector { + border-radius: 8px; + pointer-events: none !important; + padding-right: 2px; +} + +.product-dropdown-trigger svg{ + display: none !important; +} + +.fern-product-selector .product-dropdown-trigger p{ + font-weight: bold !important; +} +.fern-product-selector-radio-group { + display: grid; + grid-template-columns: repeat(3, 1fr); + gap: 8px; + max-width: 1000px; +} + +@media (max-width: 768px) { + .fern-product-selector-radio-group { + grid-template-columns: repeat(2, 1fr); + } +} +.fern-version-selector { + transform: translateY(-1px); +} + +.fern-version-selector .version-dropdown-trigger{ + outline: 1px solid var(--border, var(--grayscale-a5)) !important; + border-radius: 5px; + transition: box-shadow 0.3s ease, outline 0.3s ease; +} +.product-dropdown-trigger{ + padding-left: 0px !important; +} + +.product-dropdown-trigger, .version-dropdown-trigger{ + background-color: transparent !important; +} +.product-dropdown-trigger svg:hover{ + stroke: var(--nv-color-green) !important; +} +.version-dropdown-trigger:hover{ + box-shadow: 0 0 0 1px var(--nv-color-green) !important; +} +.version-dropdown-trigger svg:hover{ + stroke: var(--nv-color-green) !important; +} +/* Sidebar styling */ +#fern-sidebar { + border-right: 1px solid var(--border, var(--grayscale-a5)) !important; + height: 100vh !important; +} +.fern-sidebar-link:not(:hover){ + background-color: transparent !important; +} +.fern-sidebar-link { + padding-left: 1rem !important; + padding-right: 1rem !important; + padding-top: 0.5rem !important; + padding-bottom: 0.5rem !important; + border-radius: 0px !important; + &.nested { + padding-left: 1rem !important; + } +} +/* Section-level sidebar links (pages that have children) should match sidebar heading padding */ +.fern-sidebar-group > li > .fern-sidebar-link:has(+ .fern-sidebar-group) { + padding-left: 0.25rem !important; +} +.fern-sidebar-group{ + padding: 0 !important +} +#fern-sidebar-scroll-area{ + padding-right: 0 !important +} + +/* header styling */ +.fern-header-content{ + padding-left: 18.5px; + margin-top: -5px; + margin-bottom: -5px; +} +#fern-header { + border-color: var(--border, var(--grayscale-a5)) !important; +} +@keyframes header-background-fade { + 0% { + background-color: transparent; + } + 100% { + background-color: var(--header-background); + } + } + +[data-theme=default]#fern-header { +animation: header-background-fade linear; +animation-timeline: scroll(); +animation-range: 0 50px; +} +.fern-header-navbar-links .fern-button{ + background-color: transparent !important; +} +.fern-header-navbar-links > button{ + background-color: transparent !important; +} +.fern-header-logo-container > div > div > a > img{ + padding-right: 0.5rem; +} +.fern-header-logo-container .font-heading{ + font-size: 16px !important; + font-weight: bold !important; + color: var(--grayscale-a12) !important; + border-inline: 1px solid var(--border, var(--grayscale-a5)); + padding: 15px 1rem; + margin: -20px 0.5rem; +} +@media (max-width: 1024px) { + .fern-header-logo-container .font-heading{ + display: none !important; + } +} +/* Search bar styling */ +#fern-search-button{ + background-color: transparent !important; + border-radius: var(--rounded); + transition: box-shadow 0.3s ease, outline 0.3s ease; +} +#fern-search-button:hover{ + box-shadow: 0 0 0 1px var(--nv-color-green) !important; +} +#fern-search-button .fern-kbd{ + display: none; +} + +.fern-layout-footer-toolbar button{ + background-color: transparent !important; + border-color: transparent !important; + padding-inline: 0px !important; +} + +/* ========== Custom footer (native React component) – 1:1 with original ========== */ +.bd-footer { + border-top: 1px solid var(--border, var(--grayscale-a5)) !important; + font-family: NVIDIA, -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif !important; + font-size: 0.875rem; + padding: 2rem 0; + width: 100%; +} +.bd-footer * { + font-family: inherit; +} +.bd-footer__inner { + padding: 0 2rem; +} +.footer-items__start { + display: flex; + flex-direction: column; + gap: 1.5rem; +} +.footer-logos-container { + display: flex; + align-items: center; + justify-content: space-between; + width: 100%; + gap: 1rem; +} +.footer-brand { + display: inline-block; + text-decoration: none; +} +.footer-brand .logo__image { + height: 24px; + width: auto; + transition: opacity 0.2s ease; +} +.footer-brand:hover .logo__image { + opacity: 0.8; +} +.footer-brand-fern { + display: flex; + align-items: center; + margin-left: auto; +} +/* Logo theme visibility – .dark is on ancestor in Fern */ +.only-light { + display: block; + filter: invert(1); +} +.only-dark { + display: none; +} +.dark .only-light { + display: none; +} +.dark .only-dark { + display: block; + filter: none; +} +.footer-links { + display: flex; + flex-wrap: wrap; + gap: 0.25rem 0.5rem; + line-height: 1.65; + margin: 0; + padding: 0; +} +.footer-links a { + color: var(--grayscale-a11); + text-decoration: none; + transition: color 0.2s ease; + white-space: nowrap; +} +.pipe-separator { + color: var(--grayscale-a11); + white-space: nowrap; +} +.copyright { + color: var(--grayscale-a11); + font-size: 0.875rem; + line-height: 1.65; + margin: 0; +} +@media (max-width: 768px) { + .bd-footer { padding: 1.5rem 0; } + .bd-footer__inner { padding: 0 1.5rem; } + .footer-items__start { gap: 1rem; } + .footer-links { flex-direction: row; gap: 0.5rem 0.75rem; } + .footer-links a { white-space: normal; word-break: break-word; } +} +@media (max-width: 480px) { + .footer-links { gap: 0.5rem; } + .footer-links a { font-size: 0.8125rem; } + .copyright { font-size: 0.8125rem; } +} +/* Built with Fern link + tooltip */ +.built-with-fern-link { + display: flex; + align-items: baseline; + gap: 0.25rem; + text-decoration: none; + position: relative; +} +.built-with-fern-logo { + height: 1rem; + margin: 0; + transition: filter 150ms ease; +} +.built-with-fern-logo path { fill: var(--grayscale-a12); } +.built-with-fern-link:hover .built-with-fern-logo { filter: saturate(1) opacity(1); } +.built-with-fern-link:hover .built-with-fern-logo path:nth-child(2) { fill: #51C233; } +.built-with-fern-tooltip { + position: absolute; + top: 50%; + right: calc(100%); + bottom: auto; + left: auto; + transform: translateY(-50%); + margin: 0; + margin-right: 0.5rem; + padding: 0.5rem 0.75rem; + background-color: #FFFFFF; + color: #000000; + font-size: 0.85rem; + border-radius: 0.375rem; + border: 1px solid var(--grayscale-a5); + white-space: nowrap; + pointer-events: none; + opacity: 0; + transition: opacity 150ms ease; + transition-delay: 0s; + z-index: 50; + box-shadow: 0 4px 12px rgba(0, 0, 0, 0.15); + width: max-content; +} +.built-with-fern-link:hover .built-with-fern-tooltip { + opacity: 1; + transition-delay: 0.75s; +} +.dark .built-with-fern-tooltip { + background-color: #000000; + color: #FFFFFF; +} +.built-with-fern-logo-dark { display: none; } +.dark .built-with-fern-logo-light { display: none; } +.dark .built-with-fern-logo-dark { display: block; } +@media (prefers-color-scheme: dark) { + .built-with-fern-logo-light { display: none; } + .built-with-fern-logo-dark { display: block; } +} + +/* Footer styling */ +.fern-footer-nav{ + border-radius: var(--rounded); + background-color: transparent !important; + transition: box-shadow 0.3s ease, outline 0.3s ease; +} +/* Hide line numbers */ +.code-block-line-gutter { + display: none !important; +} +.fern-footer-prev h4, .fern-footer-next h4{ + font-size: inherit !important; +} +.fern-sidebar-link.nested[data-state="active"]:before { + left: -0px !important; + bottom: -0px !important; + top: -0px !important; + width: 2px !important; +} +.fern-sidebar-link[data-state="active"] { + color: unset !important; +} + +.fern-selection-item .fern-selection-item-icon{ + border-color: transparent !important; +} +/* Button styling */ +.fern-button{ + border-radius: var(--rounded); + font-weight: bold; +} +.fern-button.filled.primary{ + color: var(--nv-color-black); +} +.dark .fern-button.filled.primary{ + background-color: var(--nv-color-white); +} +.dark .fern-button.filled.primary:hover{ + background-color: var(--nv-light-grey-2); +} +.fern-button.outlined.normal{ + background-color: transparent; + --tw-ring-color: transparent; + color: var(--nv-color-black); +} +.fern-button.outlined.normal:hover{ + color: var(--nv-color-green) +} +.dark .fern-button.outlined.normal{ + color: var(--nv-color-white); +} +.dark .fern-button.outlined.normal:hover{ + color: var(--nv-color-green); +} +/* Card styling */ +.fern-card{ + transition: box-shadow 0.3s ease, outline 0.3s ease; +} +svg.card-icon{ + height: 24px !important; + width: 24px !important; +} +.card-icon{ + background-color: transparent !important; +} +.fern-card:hover{ + box-shadow: 0 0 0 1px var(--nv-color-green) !important; +} +.fern-docs-badge{ + border-radius: var(--rounded); +} +.fern-page-actions button:hover{ + background-color: transparent !important; +} +.fern-page-actions a:hover{ + background-color: transparent !important; +} +/* Moving logo to footer */ +#builtwithfern, #builtwithfern * { + display: none !important; +} + +/* Landing Page Gradients */ +/* Top: Simple radial gradient (no mask, responsive) */ +.landing-gradient-top { + position: absolute; + top: 0; + left: 0; + right: 0; + height: 800px; + background: radial-gradient(ellipse 100% 100% at 50% 10%, + rgba(191, 242, 48, 0.15) 0%, + rgba(158, 228, 179, 0.12) 30%, + rgba(124, 215, 254, 0.12) 50%, + rgba(124, 215, 254, 0.06) 75%, + transparent 100%); + pointer-events: none; + z-index: 0; +} + +/* Bottom: Masked gradient for organic transition */ +.landing-gradient-bottom { + position: absolute; + bottom: -282px; + left: 0; + right: 0; + height: 1232px; + background: linear-gradient(85deg, #BFF230 41.98%, #7CD7FE 99.52%); + opacity: 0.05; + pointer-events: none; + z-index: 5; + mask-image: url('https://www.figma.com/api/mcp/asset/27509afa-9c16-46bb-8415-4395e2e5a347'); + mask-repeat: no-repeat; + mask-position: 0% -17px; + mask-size: 100% auto; + -webkit-mask-image: url('https://www.figma.com/api/mcp/asset/27509afa-9c16-46bb-8415-4395e2e5a347'); + -webkit-mask-repeat: no-repeat; + -webkit-mask-position: 0% -17px; + -webkit-mask-size: 100% auto; +} + +/* Landing Page Gradients Wrapper */ +.landing-page-gradients { + position: relative; + width: 100%; + margin-top: -100px; + padding-top: 100px; + overflow: visible; + background: #181818; +} + +/* Hero Section (Landing page only) */ +.hero-section { + position: relative; + width: 100%; + padding: 3rem 6rem; + margin: 0 auto; + overflow: visible; + display: flex; + flex-direction: column; + align-items: center; + z-index: 10; +} + +/* Hero Section Content - constrain width */ +.hero-section > * { + position: relative; + z-index: 100; + max-width: 1440px; + width: 100%; +} + +/* Tablet and Mobile: fix spacing and layout */ +@media (max-width: 1024px) { + /* Extend dark background behind header */ + .landing-page body, .landing-page html, .landing-page main { + background: #181818 !important; + } + + .landing-page-gradients { + margin-top: -100px; + padding-top: 100px; + } + + .hero-section { + padding: 2rem 2rem; + } + + .hero-section > * { + max-width: none; + } + + .hero-content-grid { + grid-template-columns: 1fr; + gap: 2rem; + } + + .hero-heading { + font-size: 36px; + } + + .hero-subtitle { + font-size: 16px; + } + + .hero-title-section { + margin-bottom: 2rem; + } +} + +/* Small mobile only */ +@media (max-width: 600px) { + .hero-heading { + font-size: 28px; + } + + .hero-section { + padding: 1.5rem 1.5rem; + } +} + +.hero-section h1, +.hero-section h2, +.hero-section h3, +.hero-section h4, +.hero-section h5, +.hero-section h6 { + pointer-events: none !important; +} +/* Hero Title Section */ +.hero-title-section { + text-align: center; + margin-bottom: 4rem; + position: relative; + z-index: 100; +} + +.hero-heading { + font-size: 48px; + font-weight: 700; + line-height: 1.2; + margin: 0 0 1rem 0; + color: var(--nv-color-white); +} + +.hero-subtitle { + font-size: 18px; + line-height: 1.5; + margin: 0; + color: var(--nv-color-white); +} + +/* Hero Content Grid */ +.hero-content-grid { + display: grid; + grid-template-columns: repeat(2, 1fr); + gap: 3rem; + align-items: start; + position: relative; + z-index: 100; +} + +.hero-column { + display: flex; + flex-direction: column; + gap: 1rem; +} + +.hero-column-title { + font-size: 24px; + font-weight: 700; + margin: 0; + color: var(--nv-color-white); +} + +.hero-column-subtitle { + font-size: 16px; + margin: 0 0 1rem 0; + color: var(--nv-color-white); +} + +/* Hero Card Container (Left Column) */ +.hero-card-container { + display: flex; + flex-direction: column; + border-radius: 8px; + overflow: hidden; + border: 1px solid var(--border, var(--grayscale-a5)); + margin-top: 1.5rem !important; + background: rgba(26, 26, 26, 0.2); + backdrop-filter: blur(6px); +} + +.hero-card-image { + width: 100%; + height: auto; + display: block; +} + +.hero-card-content { + padding: 1.5rem; + display: flex; + flex-direction: row; + gap: 1rem; + align-items: center; + justify-content: space-between; + background: rgba(26, 26, 26, 0.2); + backdrop-filter: blur(6px); +} + +.hero-card-text-wrapper { + flex: 1; +} + +.hero-card-text { + margin: 0; + font-size: 14px; + line-height: 1.5; + color: var(--nv-color-white); +} + +.hero-card-button-wrapper { + flex-shrink: 0; +} +.hero-card-button-wrapper .fern-mdx-link{ + text-decoration: none !important; +} + +.hero-card-button { + white-space: nowrap; +} + +/* Hero Cards */ + +.hero-column .fern-card { + padding: 9px 17px; + background-color: rgba(26, 26, 26, 0.2) !important; + backdrop-filter: blur(6px); +} + +.hero-section .fern-card{ + color: white !important; +} + +.hero-column .card-icon { + font-size: 64px !important; + width: 64px !important; + height: 64px !important; +} + +.hero-column .card-icon svg, +.hero-column .card-icon i { + font-size: 64px !important; + width: 64px !important; + height: 64px !important; +} + +.hero-column .fern-card-title { + font-size: 16px; + font-weight: 500; + line-height: 24px; +} + +.hero-column .fern-card p { + font-size: 14px; + line-height: 20px; + color: white !important; +} + +/* Body Section */ +.body-section { + display: flex; + padding: 4rem 16rem; + flex-direction: column; + justify-content: center; + align-items: center; + gap: 4rem; + align-self: stretch; + position: relative; + z-index: 1; + background: #181818; +} + +/* Body Section Content - constrain width */ +.body-section > * { + max-width: 1440px; + width: 100%; + position: relative; + z-index: 10; +} + +.code-block .fern-code-link{ + text-decoration: underline !important; + text-decoration-color: var(--accent) !important; + text-underline-offset: 1px !important; + text-decoration-style: underline !important; +} + +/* Mobile Styles */ +@media (max-width: 768px) { + .hero-section { + padding: 2rem 1.5rem; + } + + .hero-title-section { + margin-bottom: 2rem; + } + + .hero-heading { + font-size: 32px; + } + + .hero-subtitle { + font-size: 16px; + } + + .hero-content-grid { + grid-template-columns: 1fr; + gap: 2rem; + } + + .hero-column-title { + font-size: 20px; + } + + .hero-column-subtitle { + font-size: 14px; + } + + .hero-card-content { + flex-direction: column; + align-items: flex-start; + } + + .hero-card-button-wrapper { + align-self: flex-start; + } + + .hero-column .card-icon, + .hero-column .card-icon svg, + .hero-column .card-icon i { + font-size: 40px !important; + width: 40px !important; + height: 40px !important; + } + + .hero-column .fern-card-title { + font-size: 14px; + } + + .hero-column .fern-card p { + font-size: 11px; + } + + .body-section { + padding: 2rem 1.5rem; + } + + .fern-selection-item-icon.use-icon { + display: none !important; + } +} \ No newline at end of file diff --git a/fern/scripts/add_frontmatter.py b/fern/scripts/add_frontmatter.py new file mode 100644 index 0000000000..6fc78fb01c --- /dev/null +++ b/fern/scripts/add_frontmatter.py @@ -0,0 +1,67 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Add frontmatter (title, description) to MDX files derived from first H1.""" + +import argparse +import re +from pathlib import Path + + +def derive_title(content: str) -> str: + """Extract title from first # Heading.""" + match = re.search(r"^#\s+(.+)$", content, re.MULTILINE) + if match: + title = match.group(1).strip() + title = re.sub(r"\{[^}]+\}`[^`]*`", "", title).strip() + return title or "Untitled" + return "Untitled" + + +def add_frontmatter(filepath: Path) -> bool: + """Add frontmatter if missing. Returns True if changes were made.""" + content = filepath.read_text() + + if content.strip().startswith("---"): + return False + + title = derive_title(content) + title_escaped = title.replace('"', '\\"') + frontmatter = f'---\ntitle: "{title_escaped}"\ndescription: ""\n---\n\n' + body = content.lstrip() + + # Remove duplicate H1 that matches title (Fern uses frontmatter title) + body = re.sub(r"^#\s+" + re.escape(title) + r"\s*\n+", "", body, count=1) + + new_content = frontmatter + body + filepath.write_text(new_content) + return True + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Add frontmatter to MDX files" + ) + parser.add_argument( + "pages_dir", + type=Path, + help="Path to pages directory (e.g. fern/v0.2.0/pages)", + ) + args = parser.parse_args() + + pages_dir = args.pages_dir.resolve() + if not pages_dir.exists(): + raise SystemExit(f"Error: pages directory not found at {pages_dir}") + + changed = [] + for mdx_file in sorted(pages_dir.rglob("*.mdx")): + if add_frontmatter(mdx_file): + changed.append(mdx_file.relative_to(pages_dir)) + print(f" Added frontmatter: {mdx_file.relative_to(pages_dir)}") + + print(f"\nAdded frontmatter to {len(changed)} files") + + +if __name__ == "__main__": + main() diff --git a/fern/scripts/check_unconverted.sh b/fern/scripts/check_unconverted.sh new file mode 100755 index 0000000000..a2795e8ffb --- /dev/null +++ b/fern/scripts/check_unconverted.sh @@ -0,0 +1,74 @@ +#!/bin/bash +# Check for unconverted MyST syntax in Fern docs + +set -e + +PAGES_DIR="${1:-fern/v0.2.0/pages}" + +echo "=== Checking for unconverted MyST syntax in $PAGES_DIR ===" +echo "" + +ISSUES_FOUND=0 + +echo "Checking for MyST directives (:::)..." +if grep -r ':::' "$PAGES_DIR" 2>/dev/null; then + echo "⚠️ Found unconverted MyST directives (see above)" + ISSUES_FOUND=1 +else + echo "✓ No MyST directives found" +fi +echo "" + +echo "Checking for {ref} references (Sphinx cross-refs, not LaTeX \\text{ref})..." +if grep -rE '\{ref\}`' "$PAGES_DIR" 2>/dev/null || grep -rE '\{ref\} ' "$PAGES_DIR" 2>/dev/null; then + echo "⚠️ Found unconverted {ref} references" + ISSUES_FOUND=1 +else + echo "✓ No {ref} references found" +fi +echo "" + +echo "Checking for {octicon} icons..." +if grep -r '{octicon}' "$PAGES_DIR" 2>/dev/null; then + echo "⚠️ Found unconverted {octicon} icons" + ISSUES_FOUND=1 +else + echo "✓ No {octicon} icons found" +fi +echo "" + +echo "Checking for {py:class} / {py:meth} / {py:mod} / {py:attr} / {py:func} / {doc}..." +if grep -rE '\{py:(class|meth|mod|attr|func)\}' "$PAGES_DIR" 2>/dev/null || grep -rE '\{doc\}`' "$PAGES_DIR" 2>/dev/null; then + echo "⚠️ Found unconverted py: or doc: roles" + ISSUES_FOUND=1 +else + echo "✓ No py:/doc roles found" +fi +echo "" + +echo "Checking for sphinx-design badges..." +if grep -r '{bdg-' "$PAGES_DIR" 2>/dev/null; then + echo "⚠️ Found unconverted badges" + ISSUES_FOUND=1 +else + echo "✓ No badges found" +fi +echo "" + +echo "Checking for MyST mermaid syntax..." +if grep -r '```{mermaid}' "$PAGES_DIR" 2>/dev/null; then + echo "⚠️ Found unconverted mermaid blocks (should be \`\`\`mermaid)" + ISSUES_FOUND=1 +else + echo "✓ No MyST mermaid syntax found" +fi +echo "" + +echo "=== Summary ===" +if [ $ISSUES_FOUND -eq 0 ]; then + echo "✓ All checks passed" + exit 0 +else + echo "⚠️ Some issues found - review and fix above" + exit 1 +fi diff --git a/fern/scripts/convert_mb_specific.py b/fern/scripts/convert_mb_specific.py new file mode 100644 index 0000000000..681d4cdd39 --- /dev/null +++ b/fern/scripts/convert_mb_specific.py @@ -0,0 +1,97 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Convert Megatron-Bridge-specific MyST/Sphinx syntax: {py:*}, {doc}.""" + +import argparse +import re +from pathlib import Path + +API_DOCS_BASE = "https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs" + + +def escape_mdx_curly_braces(content: str) -> str: + """Escape {variable} in code blocks so MDX doesn't parse as JSX.""" + return content.replace("{overrides}", "\\{overrides\\}") + + +def py_ref_to_api_url(ref: str) -> str: + """Convert dotted Python ref to API docs URL. e.g. bridge.models -> .../bridge/bridge.models.html""" + ref = ref.strip().lstrip("~") + # API structure: apidocs/bridge/.html + return f"{API_DOCS_BASE}/bridge/{ref}.html" + + +def replace_py_role(match: re.Match[str]) -> str: + """Convert py role to link to API docs (or inline code if not a bridge ref).""" + text = match.group(1).strip().lstrip("~") + # Only link bridge.* and megatron.* refs; others stay as inline code + if text.startswith("bridge.") or text.startswith("megatron."): + url = py_ref_to_api_url(text) + return f"[`{text}`]({url})" + return f"`{text}`" + + +def convert_py_roles(content: str) -> str: + """Convert {py:class}, {py:meth}, {py:mod}, {py:attr}, {py:func} to links or inline code.""" + pattern = r"\{py:(?:class|meth|mod|attr|func)\}`([^`<]+?)(?:\s*<[^>]+>)?`" + return re.sub(pattern, replace_py_role, content) + + +def convert_doc_roles(content: str) -> str: + """Convert {doc}`path` to internal links. apidocs paths -> API docs URL.""" + def replace_doc(match: re.Match[str]) -> str: + path = match.group(1).strip() + if path.startswith("apidocs/"): + return f"[API Documentation]({API_DOCS_BASE}/)" + clean = path.replace("../", "").replace(".md", "").replace(".mdx", "") + if not clean.startswith("/"): + clean = "/" + clean + display = path.split("/")[-1].replace("-", " ").replace("_", " ").title() + return f"[{display}]({clean})" + + return re.sub(r"\{doc\}`([^`<]+?)(?:\s*<[^>]+>)?`", replace_doc, content) + + +def convert_file(filepath: Path) -> bool: + """Convert a single file. Returns True if changes were made.""" + content = filepath.read_text() + original = content + + content = escape_mdx_curly_braces(content) + content = convert_py_roles(content) + content = convert_doc_roles(content) + + if content != original: + filepath.write_text(content) + return True + return False + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Convert Megatron-Bridge-specific syntax (py:class, py:meth)" + ) + parser.add_argument( + "pages_dir", + type=Path, + help="Path to pages directory (e.g. fern/v0.2.0/pages)", + ) + args = parser.parse_args() + + pages_dir = args.pages_dir.resolve() + if not pages_dir.exists(): + raise SystemExit(f"Error: pages directory not found at {pages_dir}") + + changed = [] + for mdx_file in sorted(pages_dir.rglob("*.mdx")): + if convert_file(mdx_file): + changed.append(mdx_file.relative_to(pages_dir)) + print(f" Converted: {mdx_file.relative_to(pages_dir)}") + + print(f"\nConverted {len(changed)} files") + + +if __name__ == "__main__": + main() diff --git a/fern/scripts/convert_myst_to_fern.py b/fern/scripts/convert_myst_to_fern.py new file mode 100644 index 0000000000..c2a33bb396 --- /dev/null +++ b/fern/scripts/convert_myst_to_fern.py @@ -0,0 +1,389 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Convert MyST Markdown syntax to Fern MDX components. + +Handles: admonitions, dropdowns, tab sets, grid cards, toctree removal, +HTML comments, plus Megatron-Bridge-specific: {image}, {contents}, +{literalinclude}, {admonition}, {code-block}, {doctest}. +Also converts and to Markdown links so MDX doesn't parse them as JSX. +Run convert_mb_specific.py first to strip {py:class} and {py:meth} roles. +""" + +import argparse +import re +from pathlib import Path + + +def convert_admonitions(content: str) -> str: + """Convert MyST admonitions to Fern components.""" + admonition_map = { + "note": "Note", + "warning": "Warning", + "tip": "Tip", + "important": "Info", + "seealso": "Note", + "caution": "Warning", + "danger": "Warning", + "attention": "Warning", + "hint": "Tip", + } + + for myst_type, fern_component in admonition_map.items(): + pattern = rf"```\{{{myst_type}\}}\s*\n(.*?)```" + replacement = rf"<{fern_component}>\n\1" + content = re.sub(pattern, replacement, content, flags=re.DOTALL | re.IGNORECASE) + + pattern = rf":::\{{{myst_type}\}}\s*\n(.*?):::" + content = re.sub(pattern, replacement, content, flags=re.DOTALL | re.IGNORECASE) + + return content + + +def convert_admonition_directive(content: str) -> str: + """Convert {admonition} Title :class: dropdown to Accordion.""" + pattern = r"```\{admonition\}\s+([^\n]+)(?:\s*\n(?::[^\n]+\n)*)?\n(.*?)```" + def replace(match: re.Match[str]) -> str: + title = match.group(1).strip().replace('"', "'") + body = match.group(2).strip() + return f'\n{body}\n' + return re.sub(pattern, replace, content, flags=re.DOTALL) + + +def convert_dropdowns(content: str) -> str: + """Convert MyST dropdowns to Fern Accordion components. + + Handles both fenced ```{dropdown} and directive ::: {dropdown} formats. + """ + def replace_dropdown(match: re.Match[str]) -> str: + title = match.group(1).strip() + body = match.group(2).strip() + if '"' in title: + title = title.replace('"', "'") + return f'\n{body}\n' + + # Pattern 1: ```{dropdown} Title\ncontent\n``` + pattern_fenced = r"```\{dropdown\}\s+([^\n]+)\s*\n(.*?)```" + content = re.sub(pattern_fenced, replace_dropdown, content, flags=re.DOTALL) + + # Pattern 2: ::: {dropdown} Title\ncontent\n::: + pattern_directive = r":::\s*\{dropdown\}\s+([^\n]+)(?:\s*\n(?::[^\n]+\n)*)?\n(.*?)\n:::\s*\n" + content = re.sub(pattern_directive, lambda m: replace_dropdown(m) + "\n", content, flags=re.DOTALL) + + return content + + +def convert_tab_sets(content: str) -> str: + """Convert MyST tab sets to Fern Tabs components.""" + content = re.sub(r"::::+\s*\{tab-set\}\s*", "\n", content) + content = re.sub(r"```\{tab-set\}\s*", "\n", content) + + def replace_tab_item(match: re.Match[str]) -> str: + title = match.group(1).strip() + return f'' + + content = re.sub(r"::::*\s*\{tab-item\}\s+([^\n]+)", replace_tab_item, content) + content = re.sub(r":::*\s*\{tab-item\}\s+([^\n]+)", replace_tab_item, content) + + lines = content.split("\n") + result = [] + in_tab = False + + for line in lines: + if '\n") + in_tab = True + result.append(line) + elif line.strip() in [":::::", "::::", ":::", ""]: + if in_tab and line.strip() != "": + result.append("") + in_tab = False + if line.strip() in [":::::", "::::"]: + result.append("") + else: + result.append(line) + else: + result.append(line) + + content = "\n".join(result) + content = re.sub(r"\n::::+\n", "\n", content) + content = re.sub(r"\n:::+\n", "\n", content) + return content + + +def convert_grid_cards(content: str) -> str: + """Convert MyST grid cards to Fern Cards components.""" + content = re.sub(r"::::+\s*\{grid\}[^\n]*\n", "\n", content) + content = re.sub(r"```\{grid\}[^\n]*\n", "\n", content) + + def replace_card(match: re.Match[str]) -> str: + full_match = match.group(0) + title_match = re.search(r"\{grid-item-card\}\s+(.+?)(?:\n|$)", full_match) + title = title_match.group(1).strip() if title_match else "Card" + link_match = re.search(r":link:\s*(\S+)", full_match) + href = link_match.group(1) if link_match else "" + if href and href != "apidocs/index": + if not href.startswith("http"): + href = "/" + href.replace(".md", "").replace(".mdx", "") + return f'' + if href == "apidocs/index": + return f'' + return f'' + + content = re.sub( + r"::::*\s*\{grid-item-card\}[^\n]*(?:\n:link:[^\n]*)?(?:\n:link-type:[^\n]*)?", + replace_card, + content, + ) + content = re.sub( + r":::*\s*\{grid-item-card\}[^\n]*(?:\n:link:[^\n]*)?(?:\n:link-type:[^\n]*)?", + replace_card, + content, + ) + + lines = content.split("\n") + result = [] + in_card = False + + for line in lines: + if '\n") + in_card = True + result.append(line) + elif line.strip() in [":::::", "::::", ":::", ""]: + if in_card and line.strip() != "": + result.append("\n") + in_card = False + if line.strip() in [":::::", "::::"]: + result.append("\n") + else: + result.append(line) + + return "\n".join(result) + + +def convert_list_table(content: str) -> str: + """Convert MyST list-table to markdown table. + + Handles ```{list-table} with * - cell format. + """ + pattern = r"```\{list-table\}[^\n]*(?:\n:[^\n]+)*\n\n(.*?)```" + + def replace_list_table(match: re.Match[str]) -> str: + body = match.group(1).strip() + rows: list[list[str]] = [] + for line in body.split("\n"): + line = line.rstrip() + if not line: + continue + if line.startswith("* -"): + rows.append([line[3:].strip()]) + elif line.startswith(" -") or line.startswith("-"): + cell = line.lstrip("- ").strip() + if rows: + rows[-1].append(cell) + else: + rows.append([cell]) + if not rows: + return match.group(0) + header = rows[0] + sep = "| " + " | ".join(["---"] * len(header)) + " |" + lines_out = ["| " + " | ".join(header) + " |", sep] + for row in rows[1:]: + while len(row) < len(header): + row.append("") + lines_out.append("| " + " | ".join(row[: len(header)]) + " |") + return "\n".join(lines_out) + + return re.sub(pattern, replace_list_table, content, flags=re.DOTALL) + + +def remove_toctree(content: str) -> str: + """Remove toctree blocks entirely.""" + content = re.sub(r"```\{toctree\}.*?```", "", content, flags=re.DOTALL) + content = re.sub(r":::\{toctree\}.*?:::", "", content, flags=re.DOTALL) + return content + + +def remove_contents(content: str) -> str: + """Remove {contents} directive (Fern has its own nav).""" + content = re.sub(r"```\{contents\}.*?```", "", content, flags=re.DOTALL) + content = re.sub(r":::\{contents\}.*?:::", "", content, flags=re.DOTALL) + return content + + +def convert_image(content: str, filepath: Path, repo_root: Path) -> str: + """Convert {image} path to markdown image. Path relative to current file.""" + pattern = r"```\{image\}\s+([^\s\n]+)(?:\s*\n(?::[^\n]+\n)*)?```" + def replace(match: re.Match[str]) -> str: + img_path = match.group(1).strip() + # Path is relative to current file's dir in docs/; we copy to fern/assets/training/images/ + # output path for update_links: /assets/training/images/filename + if "images/" in img_path: + img_name = img_path.split("images/")[-1] + return f"![{img_name}](/assets/training/images/{img_name})" + return f"![{img_path}]({img_path})" + return re.sub(pattern, replace, content) + + +def convert_literalinclude(content: str, filepath: Path, repo_root: Path) -> str: + """Convert {literalinclude} to fenced code block. Inlines full file.""" + pattern = r"```\{literalinclude\}\s+([^\s\n]+)(?:\s*\n(?::[^\n]+\n)*)?\s*```" + def replace(match: re.Match[str]) -> str: + inc_path = match.group(1).strip() + # Path is relative to docs/ (e.g. ../src/megatron/...) + resolved = (repo_root / "docs" / inc_path).resolve() + if not resolved.exists(): + resolved = (repo_root / inc_path.replace("../", "")).resolve() + if not resolved.exists(): + return f"" + lang = "python" if resolved.suffix == ".py" else "" + try: + body = resolved.read_text() + except Exception: + return f"" + return f"```{lang}\n{body}\n```" + return re.sub(pattern, replace, content) + + +def convert_code_block(content: str) -> str: + """Convert {code-block} lang to standard ```lang.""" + pattern = r"```\{code-block\}\s+(\w+)(?:\s*\n(?::[^\n]+\n)*)?\n(.*?)```" + def replace(match: re.Match[str]) -> str: + lang = match.group(1) + body = match.group(2).rstrip() + return f"```{lang}\n{body}\n```" + return re.sub(pattern, replace, content, flags=re.DOTALL) + + +def convert_doctest(content: str) -> str: + """Convert {doctest} to standard code block.""" + pattern = r"```\{doctest\}\s*\n(.*?)```" + def replace(match: re.Match[str]) -> str: + body = match.group(1).strip() + return f"```python\n{body}\n```" + return re.sub(pattern, replace, content, flags=re.DOTALL) + + +def escape_sphinx_doc_refs(content: str) -> str: + """Escape Sphinx doc refs like that MDX parses as JSX.""" + content = re.sub( + r"", + "[API Documentation](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/)", + content, + ) + return content + + +def convert_angle_bracket_urls_and_emails(content: str) -> str: + """Convert and to Markdown links so MDX doesn't parse them as JSX tags.""" + # or -> [url](url) + content = re.sub( + r"<(https?://[^>]+)>", + r"[\1](\1)", + content, + ) + # -> [email](mailto:email) + content = re.sub( + r"<([a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,})>", + r"[\1](mailto:\1)", + content, + ) + return content + + +def convert_html_comments(content: str) -> str: + """Convert HTML comments to JSX comments.""" + return re.sub(r"", r"{/* \1 */}", content, flags=re.DOTALL) + + +def remove_directive_options(content: str) -> str: + """Remove MyST directive options.""" + for opt in [ + ":icon:", ":class:", ":columns:", ":gutter:", ":margin:", ":padding:", + ":link-type:", ":maxdepth:", ":titlesonly:", ":hidden:", ":link:", + ":caption:", ":language:", ":pyobject:", ":linenos:", ":emphasize-lines:", + ":width:", ":align:", ":relative-docs:", + ]: + content = re.sub(rf"\n{re.escape(opt)}[^\n]*", "", content) + return content + + +def fix_malformed_tags(content: str) -> str: + """Fix common malformed tag issues.""" + content = re.sub(r'title=""', 'title="Details"', content) + content = re.sub( + r"<(Note|Warning|Tip|Info)([^>]*)/>\s*\n([^<]+)", + r"<\1\2>\n\3", + content, + ) + return content + + +def clean_multiple_newlines(content: str) -> str: + """Clean up excessive newlines.""" + content = re.sub(r"\n{3,}", "\n\n", content) + return content.strip() + "\n" + + +def convert_file(filepath: Path, repo_root: Path) -> bool: + """Convert a single file. Returns True if changes were made.""" + content = filepath.read_text() + original = content + + content = convert_admonitions(content) + content = convert_admonition_directive(content) + content = convert_dropdowns(content) + content = convert_grid_cards(content) + content = convert_tab_sets(content) + content = convert_list_table(content) + content = remove_toctree(content) + content = remove_contents(content) + content = convert_image(content, filepath, repo_root) + content = convert_literalinclude(content, filepath, repo_root) + content = convert_code_block(content) + content = convert_doctest(content) + content = escape_sphinx_doc_refs(content) + content = convert_angle_bracket_urls_and_emails(content) + content = convert_html_comments(content) + content = remove_directive_options(content) + content = fix_malformed_tags(content) + content = clean_multiple_newlines(content) + + if content != original: + filepath.write_text(content) + return True + return False + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Convert MyST syntax to Fern MDX in pages directory" + ) + parser.add_argument( + "pages_dir", + type=Path, + help="Path to pages directory (e.g. fern/v0.2.0/pages)", + ) + args = parser.parse_args() + + pages_dir = args.pages_dir.resolve() + if not pages_dir.exists(): + raise SystemExit(f"Error: pages directory not found at {pages_dir}") + + repo_root = pages_dir.parent.parent.parent + + changed = [] + for mdx_file in sorted(pages_dir.rglob("*.mdx")): + if convert_file(mdx_file, repo_root): + changed.append(mdx_file.relative_to(pages_dir)) + print(f" Converted: {mdx_file.relative_to(pages_dir)}") + + print(f"\nConverted {len(changed)} files") + + +if __name__ == "__main__": + main() diff --git a/fern/scripts/copy_docs_to_fern.py b/fern/scripts/copy_docs_to_fern.py new file mode 100644 index 0000000000..72e5622734 --- /dev/null +++ b/fern/scripts/copy_docs_to_fern.py @@ -0,0 +1,108 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Copy docs/*.md to fern//pages/*.mdx preserving directory structure.""" + +import argparse +import shutil +from pathlib import Path + +SKIP_FILES = { + "conf.py", + "Makefile", + "helpers.py", + "versions1.json", + "project.json", +} +SKIP_DIRS = {"_templates", "_build", "apidocs", ".venv", ".git"} + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Copy docs/*.md to fern//pages/*.mdx" + ) + parser.add_argument( + "version", + help="Version folder name (e.g. v0.2.0)", + ) + parser.add_argument( + "--docs-dir", + default="docs", + help="Source docs directory (default: docs)", + ) + parser.add_argument( + "--fern-dir", + default="fern", + help="Fern root directory (default: fern)", + ) + args = parser.parse_args() + + repo_root = Path(__file__).resolve().parent.parent.parent + docs_dir = repo_root / args.docs_dir + fern_dir = repo_root / args.fern_dir + pages_dir = fern_dir / args.version / "pages" + + if not docs_dir.exists(): + raise SystemExit(f"Error: docs directory not found at {docs_dir}") + + pages_dir.mkdir(parents=True, exist_ok=True) + + # Copy docs/assets to fern/assets if they exist + docs_assets = docs_dir / "assets" + fern_assets = fern_dir / "assets" + if docs_assets.exists(): + for asset in docs_assets.rglob("*"): + if asset.is_file(): + rel = asset.relative_to(docs_assets) + dst = fern_assets / rel + dst.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(asset, dst) + print(f"Copied assets from {docs_assets} to {fern_assets}") + + # Copy docs/training/images and docs/images to fern/assets/training/images (Megatron-Bridge) + fern_training_images = fern_assets / "training" / "images" + fern_training_images.mkdir(parents=True, exist_ok=True) + + docs_training_images = docs_dir / "training" / "images" + if docs_training_images.exists(): + for img in docs_training_images.iterdir(): + if img.is_file(): + shutil.copy2(img, fern_training_images / img.name) + print(f"Copied training images from {docs_training_images} to {fern_training_images}") + + docs_images = docs_dir / "images" + if docs_images.exists(): + for img in docs_images.iterdir(): + if img.is_file(): + shutil.copy2(img, fern_training_images / img.name) + print(f"Copied docs/images to {fern_training_images}") + + # Copy repo root images (e.g. Repo-Mbridge.png from README) to fern/assets + repo_root = docs_dir.parent + for img_name in ["Repo-Mbridge.png"]: + src = repo_root / img_name + if src.exists(): + shutil.copy2(src, fern_assets / img_name) + print(f"Copied {img_name} to {fern_assets}") + + copied = 0 + for md_file in docs_dir.rglob("*.md"): + rel = md_file.relative_to(docs_dir) + + if rel.name in SKIP_FILES: + continue + if any(part in SKIP_DIRS or part.startswith(".") for part in rel.parts): + continue + + mdx_path = pages_dir / rel.with_suffix(".mdx") + mdx_path.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(md_file, mdx_path) + copied += 1 + print(f" {rel} -> {args.version}/pages/{rel.with_suffix('.mdx')}") + + print(f"\nCopied {copied} files to {pages_dir}") + + +if __name__ == "__main__": + main() diff --git a/fern/scripts/expand_includes.py b/fern/scripts/expand_includes.py new file mode 100644 index 0000000000..c1516a1292 --- /dev/null +++ b/fern/scripts/expand_includes.py @@ -0,0 +1,89 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Expand {include} directives in MDX files. Run after copy_docs_to_fern.py. + +Processes index.mdx and releases/changelog.mdx, replacing {include} blocks +with the actual content of the referenced files (README.md, CHANGELOG.md). +""" + +import argparse +import re +from pathlib import Path + + +def expand_include_in_content( + content: str, file_path: Path, pages_dir: Path, docs_dir: Path +) -> str: + """Replace {include} directives with file content. Paths are relative to the source doc.""" + # Match ```{include} path with optional options (e.g. :relative-docs:) + pattern = r"```\{include\}\s+([^\s\n]+)(?:\s*\n(?::[^\n]+\n)*)?```" + + def replace_include(match: re.Match[str]) -> str: + include_path_str = match.group(1).strip() + # Include paths are relative to the source doc's directory in docs/ + # e.g. docs/index.md has ../README.md -> repo_root/README.md + # e.g. docs/releases/changelog.md has ../../CHANGELOG.md -> repo_root/CHANGELOG.md + rel = file_path.relative_to(pages_dir) + source_dir = docs_dir / rel.parent + if rel.name == "index.mdx": + source_dir = docs_dir + resolved = (source_dir / include_path_str).resolve() + + if not resolved.exists(): + return f"" + return resolved.read_text() + + return re.sub(pattern, replace_include, content) + + +def expand_file(filepath: Path, pages_dir: Path, docs_dir: Path) -> bool: + """Expand includes in a single file. Returns True if changes were made.""" + content = filepath.read_text() + if "{include}" not in content: + return False + + new_content = expand_include_in_content(content, filepath, pages_dir, docs_dir) + if new_content != content: + filepath.write_text(new_content) + return True + return False + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Expand {include} directives in MDX files" + ) + parser.add_argument( + "pages_dir", + type=Path, + help="Path to pages directory (e.g. fern/v0.2.0/pages)", + ) + parser.add_argument( + "--docs-dir", + type=Path, + default=None, + help="Path to docs directory (default: repo_root/docs)", + ) + args = parser.parse_args() + + pages_dir = args.pages_dir.resolve() + if not pages_dir.exists(): + raise SystemExit(f"Error: pages directory not found at {pages_dir}") + + repo_root = pages_dir.parent.parent.parent + docs_dir = (args.docs_dir.resolve() if args.docs_dir else repo_root / "docs") + + expanded = [] + for pattern in ["index.mdx", "releases/changelog.mdx"]: + filepath = pages_dir / pattern + if filepath.exists() and expand_file(filepath, pages_dir, docs_dir): + expanded.append(filepath.relative_to(pages_dir)) + print(f" Expanded: {filepath.relative_to(pages_dir)}") + + print(f"\nExpanded {len(expanded)} files") + + +if __name__ == "__main__": + main() diff --git a/fern/scripts/remove_duplicate_h1.py b/fern/scripts/remove_duplicate_h1.py new file mode 100644 index 0000000000..1488122feb --- /dev/null +++ b/fern/scripts/remove_duplicate_h1.py @@ -0,0 +1,59 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Remove duplicate H1 that matches frontmatter title.""" + +import argparse +import re +from pathlib import Path + + +def remove_duplicate_h1(filepath: Path) -> bool: + """Remove H1 after frontmatter if it duplicates the title. Returns True if changed.""" + content = filepath.read_text() + + if not content.strip().startswith("---"): + return False + + # Extract title from frontmatter + match = re.search(r"^---\s*\ntitle:\s*(.+?)\n", content, re.MULTILINE) + if not match: + return False + + title = match.group(1).strip().strip('"\'') + pattern = rf"(---\s*\n.*?---\s*\n\n)#\s+{re.escape(title)}\s*\n+" + new_content = re.sub(pattern, r"\1", content, count=1, flags=re.DOTALL) + + if new_content != content: + filepath.write_text(new_content) + return True + return False + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Remove duplicate H1 that matches frontmatter title" + ) + parser.add_argument( + "pages_dir", + type=Path, + help="Path to pages directory", + ) + args = parser.parse_args() + + pages_dir = args.pages_dir.resolve() + if not pages_dir.exists(): + raise SystemExit(f"Error: pages directory not found at {pages_dir}") + + changed = [] + for mdx_file in sorted(pages_dir.rglob("*.mdx")): + if remove_duplicate_h1(mdx_file): + changed.append(mdx_file.relative_to(pages_dir)) + print(f" Removed H1: {mdx_file.relative_to(pages_dir)}") + + print(f"\nRemoved duplicate H1 from {len(changed)} files") + + +if __name__ == "__main__": + main() diff --git a/fern/scripts/update_links.py b/fern/scripts/update_links.py new file mode 100644 index 0000000000..0b823e5e4b --- /dev/null +++ b/fern/scripts/update_links.py @@ -0,0 +1,88 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Update internal links: .md -> Fern paths, relative paths -> absolute.""" + +import argparse +import re +from pathlib import Path + + +def normalize_url(url: str) -> str: + """Normalize a URL to Fern path format.""" + clean = url.replace(".md", "").replace(".mdx", "") + if url.startswith(("http://", "https://", "#", "mailto:")): + return url + # Normalize asset paths to /assets/ + if "assets/" in clean or clean.startswith("./assets") or clean.startswith("../assets"): + clean = "/assets/" + clean.split("assets/")[-1] + # Normalize images/ to /assets/training/images/ (for training docs) + elif "images/" in clean or clean.startswith("./images") or clean.startswith("../images"): + img_name = clean.split("images/")[-1] if "images/" in clean else clean.split("/")[-1] + clean = "/assets/training/images/" + img_name + # Repo root images (e.g. Repo-Mbridge.png from README) + elif clean.endswith(".png") and "/" not in clean: + clean = "/assets/" + clean + elif not clean.startswith("/"): + clean = "/" + clean + return clean + + +def update_links_in_content(content: str, file_dir: Path, pages_root: Path) -> str: + """Update markdown links and image paths: .md/.mdx -> Fern paths.""" + + def replace_link(match: re.Match[str]) -> str: + text, url = match.group(1), match.group(2) + clean = normalize_url(url) + return f"[{text}]({clean})" + + def replace_image(match: re.Match[str]) -> str: + alt, url = match.group(1), match.group(2) + clean = normalize_url(url) + return f"![{alt}]({clean})" + + # Process images first, then links (negative lookbehind avoids matching images) + content = re.sub(r"!\[([^\]]*)\]\(([^)]+)\)", replace_image, content) + content = re.sub(r"(? bool: + """Update links in a single file. Returns True if changes were made.""" + content = filepath.read_text() + file_dir = filepath.parent + new_content = update_links_in_content(content, file_dir, pages_root) + + if new_content != content: + filepath.write_text(new_content) + return True + return False + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Update internal links in MDX files" + ) + parser.add_argument( + "pages_dir", + type=Path, + help="Path to pages directory (e.g. fern/v0.2.0/pages)", + ) + args = parser.parse_args() + + pages_dir = args.pages_dir.resolve() + if not pages_dir.exists(): + raise SystemExit(f"Error: pages directory not found at {pages_dir}") + + changed = [] + for mdx_file in sorted(pages_dir.rglob("*.mdx")): + if update_file(mdx_file, pages_dir): + changed.append(mdx_file.relative_to(pages_dir)) + print(f" Updated: {mdx_file.relative_to(pages_dir)}") + + print(f"\nUpdated {len(changed)} files") + + +if __name__ == "__main__": + main() diff --git a/fern/v0.2.0/pages/README.mdx b/fern/v0.2.0/pages/README.mdx new file mode 100644 index 0000000000..755da33e21 --- /dev/null +++ b/fern/v0.2.0/pages/README.mdx @@ -0,0 +1,262 @@ +--- +title: "Megatron Bridge Documentation" +description: "" +--- + +Welcome to the Megatron Bridge documentation! This guide helps you navigate our comprehensive documentation to find exactly what you need for training, converting, and working with large language models and vision language models. + +## 🚀 Quick Start Paths + +### I want to + +**🏃‍♂️ Get started with model conversion** +→ Start with [Bridge Guide](/bridge-guide) for Hugging Face ↔ Megatron conversion + +**⚡ Understand parallelisms and performance** +→ Jump to [Parallelisms Guide](/parallelisms) and [Performance Guide](/performance-guide) + +**🚀 Start training a model** +→ See [Training Documentation](/training/README) for comprehensive training guides + +**📚 Find model documentation** +→ Browse [Supported Models](/models/llm/index) for LLMs or [Vision Language Models](/models/vlm/index) for VLMs + +**🔧 Migrate from NeMo 2 or Megatron-LM** +→ Check [NeMo 2 Migration Guide](/nemo2-migration-guide) or [Megatron-LM Migration Guide](/megatron-lm-to-megatron-bridge) + +**📊 Use training recipes** +→ Read [Recipe Usage](/recipe-usage) for pre-configured training recipes + +**🔌 Add support for a new model** +→ Refer to [Adding New Models](/adding-new-models) + +**📋 Check version information** +→ See [Releases Documentation](/releases/README) for versions, changelog, and known issues + +--- + +## 👥 Documentation by Role + +### For ML Engineers & Researchers + +- **Start here:** [Bridge Guide](/bridge-guide) → [Training Documentation](/training/README) +- **Deep dive:** [Performance Guide](/performance-guide) → [Training Optimization Guides](/training/README#optimization-and-performance) +- **Model support:** [Supported Models](/models/llm/index) → [Adding New Models](/adding-new-models) + +### For Training Engineers + +- **Start here:** [Training Documentation](/training/README) → [Configuration Container Overview](/training/config-container-overview) +- **Performance:** [Performance Guide](/performance-guide) → [Performance Summary](/performance-summary) +- **Parallelisms:** [Parallelisms Guide](/parallelisms) → [Training Optimization](/training/README#optimization-and-performance) + +### For Model Developers + +- **Start here:** [Bridge Guide](/bridge-guide) → [Bridge Tech Details](/bridge-tech-details) +- **Model support:** [Adding New Models](/adding-new-models) → [Model Documentation](/models/llm/index) +- **Integration:** [Bridge RL Integration](/bridge-rl-integration) + +### For DevOps & Platform Teams + +- **Start here:** [Releases Documentation](/releases/README) → [Software Versions](/releases/software-versions) +- **Troubleshooting:** [Known Issues](/releases/known-issues) +- **API Reference:** [API Documentation](/apidocs/index.rst) + +--- + +## 📚 Complete Documentation Index + +### Getting Started + +| Document | Purpose | When to Read | +|----------|---------|--------------| +| **[Bridge Guide](/bridge-guide)** | Hugging Face ↔ Megatron conversion guide | First time converting models | +| **[Bridge Tech Details](/bridge-tech-details)** | Technical details of the bridge system | Understanding bridge internals | +| **[Parallelisms Guide](/parallelisms)** | Data and model parallelism strategies | Setting up distributed training | +| **[Performance Summary](/performance-summary)** | Quick performance reference | Quick performance lookup | +| **[Performance Guide](/performance-guide)** | Comprehensive performance optimization | Optimizing training performance | + +### Model Support + +| Document | Purpose | When to Read | +|----------|---------|--------------| +| **[Large Language Models](/models/llm/index)** | LLM model documentation | Working with LLM models | +| **[Vision Language Models](/models/vlm/index)** | VLM model documentation | Working with VLM models | +| **[Adding New Models](/adding-new-models)** | Guide for adding model support | Extending model support | + +### Training and Customization + +| Document | Purpose | When to Read | +|----------|---------|--------------| +| **[Training Documentation](/training/README)** | Comprehensive training guides | Setting up and customizing training | +| **[Configuration Container Overview](/training/config-container-overview)** | Central training configuration | Understanding training configuration | +| **[Entry Points](/training/entry-points)** | Training entry points and execution | Understanding training flow | +| **[Training Loop Settings](/training/training-loop-settings)** | Training loop parameters | Configuring training parameters | +| **[Optimizer & Scheduler](/training/optimizer-scheduler)** | Optimization configuration | Setting up optimizers | +| **[Mixed Precision](/training/mixed-precision)** | Mixed precision training | Reducing memory usage | +| **[PEFT](/training/peft)** | Parameter-efficient fine-tuning | Fine-tuning with limited resources | +| **[Checkpointing](/training/checkpointing)** | Checkpoint management | Saving and resuming training | +| **[Logging](/training/logging)** | Logging and monitoring | Monitoring training progress | +| **[Profiling](/training/profiling)** | Performance profiling | Identifying bottlenecks | + +### Recipes and Workflows + +| Document | Purpose | When to Read | +|----------|---------|--------------| +| **[Recipe Usage](/recipe-usage)** | Using pre-configured training recipes | Quick training setup | +| **[Bridge RL Integration](/bridge-rl-integration)** | Reinforcement learning integration | RL training workflows | + +### Migration Guides + +| Document | Purpose | When to Read | +|----------|---------|--------------| +| **[NeMo 2 Migration Guide](/nemo2-migration-guide)** | Migrating from NeMo 2 | Upgrading from NeMo 2 | +| **[Megatron-LM Migration Guide](/megatron-lm-to-megatron-bridge)** | Migrating from Megatron-LM | Upgrading from Megatron-LM | + +### Reference + +| Document | Purpose | When to Read | +|----------|---------|--------------| +| **[API Documentation](/apidocs/index.rst)** | Complete API reference | Building integrations | +| **[Releases Documentation](/releases/README)** | Version history and known issues | Checking versions, troubleshooting | +| **[Documentation Guide](/documentation)** | Contributing to documentation | Contributing docs | + +--- + +## 🗺️ Common Reading Paths + +### 🆕 First-Time Users + +1. [Bridge Guide](/bridge-guide) *(10 min - understand conversion)* +2. [Parallelisms Guide](/parallelisms) *(15 min - understand distributed training)* +3. [Training Documentation](/training/README) *(choose your training path)* +4. [Recipe Usage](/recipe-usage) *(5 min - use pre-configured recipes)* + +### 🔧 Setting Up Training + +1. [Training Documentation](/training/README) *(overview of training system)* +2. [Configuration Container Overview](/training/config-container-overview) *(understand configuration)* +3. [Entry Points](/training/entry-points) *(how training starts)* +4. [Training Loop Settings](/training/training-loop-settings) *(configure parameters)* +5. [Logging](/training/logging) *(set up monitoring)* + +### ⚡ Performance Optimization + +1. [Performance Guide](/performance-guide) *(comprehensive optimization strategies)* +2. [Performance Summary](/performance-summary) *(quick reference)* +3. [Mixed Precision](/training/mixed-precision) *(reduce memory usage)* +4. [Communication Overlap](/training/communication-overlap) *(optimize distributed training)* +5. [Activation Recomputation](/training/activation-recomputation) *(reduce memory footprint)* +6. [Profiling](/training/profiling) *(identify bottlenecks)* + +### 🔄 Model Conversion Workflow + +1. [Bridge Guide](/bridge-guide) *(conversion basics)* +2. [Bridge Tech Details](/bridge-tech-details) *(technical details)* +3. [Supported Models](/models/llm/index) or [Vision Language Models](/models/vlm/index) *(model-specific guides)* +4. [Adding New Models](/adding-new-models) *(extend support)* + +### 🔧 Customization and Extension + +1. [Training Documentation](/training/README) *(training customization)* +2. [PEFT](/training/peft) *(parameter-efficient fine-tuning)* +3. [Distillation](/training/distillation) *(knowledge distillation)* +4. [Adding New Models](/adding-new-models) *(add model support)* +5. [Bridge RL Integration](/bridge-rl-integration) *(RL workflows)* + +### 📦 Migration Paths + +1. [NeMo 2 Migration Guide](/nemo2-migration-guide) *(from NeMo 2)* +2. [Megatron-LM Migration Guide](/megatron-lm-to-megatron-bridge) *(from Megatron-LM)* +3. [Training Documentation](/training/README) *(new training system)* + +--- + +## 📁 Directory Structure + +### Main Documentation + +- **Guides** - Core guides for parallelisms, performance, recipes, and migration +- **Bridge Documentation** - Hugging Face ↔ Megatron conversion guides +- **Model Documentation** - Supported model families and architectures + +### Subdirectories + +#### [models/](/models/README) + +- **[llm/](/models/llm/README)** - Large Language Model documentation + - Individual model guides (Qwen, LLaMA, Mistral, etc.) + - Conversion examples and training recipes +- **[vlm/](/models/vlm/README)** - Vision Language Model documentation + - VLM model guides (Qwen VL, Gemma VL, etc.) + - Multimodal model support + +#### [training/](/training/README) + +- **Configuration** - ConfigContainer, entry points, training loop settings +- **Optimization** - Optimizer, scheduler, mixed precision, communication overlap +- **Performance** - Attention optimizations, activation recomputation, CPU offloading +- **Monitoring** - Logging, profiling, checkpointing, resiliency +- **Advanced** - PEFT, packed sequences, distillation + +#### [releases/](/releases/README) + +- **Software Versions** - Current versions and dependencies +- **Changelog** - Release history and changes +- **Known Issues** - Bugs, limitations, and workarounds + +--- + +## 🔗 How Documents Connect + +```mermaid +graph TD + A[README.md
Start Here] --> B[Bridge Guide
Model Conversion] + A --> C[Training Docs
Training Setup] + A --> D[Models
Model Support] + + B --> E[Bridge Tech Details
Technical Deep Dive] + B --> F[Supported Models
Model-Specific Guides] + + C --> G[Config Container
Configuration] + C --> H[Performance Guide
Optimization] + C --> I[Parallelisms
Distributed Training] + + G --> J[Training Loop
Training Parameters] + G --> K[Optimizer & Scheduler
Optimization Setup] + + H --> L[Mixed Precision
Memory Efficiency] + H --> M[Communication Overlap
Performance] + + I --> N[Data Parallelism
DDP] + I --> O[Model Parallelism
TP/PP/VPP] + + D --> P[LLM Models
Language Models] + D --> Q[VLM Models
Vision Language Models] + + style A fill:#e1f5fe + style B fill:#f3e5f5 + style C fill:#e8f5e8 + style D fill:#fff3e0 + style H fill:#fce4ec + style I fill:#e0f2f1 +``` + +--- + +## 🤝 Getting Help + +- **GitHub Issues:** [Report bugs or request features](https://github.com/NVIDIA-NeMo/Megatron-Bridge/issues) +- **Documentation Issues:** Found something unclear? Let us know! +- **Community:** Join discussions and share experiences + +--- + +## 📖 Additional Resources + +- **[Examples](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/examples)** - Code examples and tutorials +- **[Contributing Guide](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/CONTRIBUTING.md)** - How to contribute to the project +- **[API Documentation](/apidocs/index.rst)** - Complete API reference + +--- + +**Ready to get started?** Choose your path above or dive into the [Bridge Guide](/bridge-guide) for model conversion! 🚀 diff --git a/fern/v0.2.0/pages/adding-new-models.mdx b/fern/v0.2.0/pages/adding-new-models.mdx new file mode 100644 index 0000000000..b0ea9ddcd0 --- /dev/null +++ b/fern/v0.2.0/pages/adding-new-models.mdx @@ -0,0 +1,301 @@ +--- +title: "Contribute a New Model to Megatron Bridge" +description: "" +--- + +This guide explains how to add support for a new 🤗 Hugging Face model (or family) to Megatron Bridge so to convert between HF ↔ Megatron-Core formats and participate in training recipes. + +Use this checklist-style flow: scaffold → provider mapping → parameter mappings → tests → validation. + +## Prerequisites + +- Familiarity with the Megatron Bridge repository structure. +- A working Python 3.10+ environment with Megatron Bridge installed (see [installation instructions](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/README.md#installation)), a container is recommended. +- Familiarity with Megatron-Core GPT-style modules and 🤗 Transformers config objects. +- Access to a small HF checkpoint for local testing. +- Read first: + - [Bridge user guide](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/docs/bridge-guide.md) + - [Technical details](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/docs/bridge-tech-details.md) + - [Model bridges overview](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/models/README.md) + +## 1) Decide the integration strategy + + Most GPT-style models (such as the Qwen and Llama families) can reuse the Megatron-Core GPT model by mapping their configuration. If the model requires custom building blocks (e.g., an attention variant, RoPE variant, or VLM modules), add a lightweight specialization similar to how 🤗 HuggingFace implements `modeling_xxx.py`. + +- **Standard GPT-style models**: Implement a `Provider` and a `Bridge`. For example, see the [Llama provider](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/models/llama/llama_provider.py) and [Llama bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/models/llama/llama_bridge.py). +- **Models with custom components**: If your model has custom operations or blocks (e.g., a unique attention mechanism), add a minimal modeling module in the same directory and reference it from the `Provider` (example forthcoming). + +## 2) Scaffold the model folder + +Create a folder under `src/megatron/bridge/models//` and add: + +- `_provider.py`: builds a `TransformerConfig`-compatible provider (or a subclass of an existing provider) and exposes `.provide_distributed_model()`. For example: [Llama provider](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/models/llama/llama_provider.py), [Qwen provider](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/models/qwen/qwen_provider.py), or [Qwen2 provider](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/models/qwen/qwen2_provider.py). +- `_bridge.py`: architecture-specific bridge that maps HF config → provider and defines parameter mappings. For example: [Llama bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/models/llama/llama_bridge.py), [Qwen3 bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/models/qwen/qwen3_bridge.py), or [Qwen2 bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/models/qwen/qwen2_bridge.py). +- Optional: `README.md` with any model quirks. For example: [Llama README](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/models/llama/README.md). + +## 3) Implement the Provider + +Your provider maps the Hugging Face config to Megatron-Core transformer config fields and lazily constructs the distributed model(s). Start from the generic GPT provider (`src/megatron/bridge/models/gpt_provider.py`) and specialize the necessary fields and flags: + +- Parallelism: `tensor_model_parallel_size`, `pipeline_model_parallel_size`, optional VPP/EP settings. +- Numerics: `fp16`, `bf16`, `params_dtype`, activation recomputation. +- Architecture quirks: RoPE base/scale, QK layernorm, tied embeddings, KV groups, max sequence length, etc. +- Optional custom modules: point to custom attention/MLP implementations using a layer spec if needed. + +Expose: +```python +provider = YourModelProvider(...) +provider.finalize() +model = provider.provide_distributed_model(wrap_with_ddp=False) +``` + +### Suggested Cursor prompt (Provider) [Expermental] +```text +You are working in the Megatron Bridge repo. Create `src/megatron/bridge/models//_provider.py`. + +Goal: Implement `YourModelProvider` that maps HF config → Megatron-Core transformer config and exposes `.provide_distributed_model()`. + +Requirements: +- Start from `src/megatron/bridge/models/gpt_provider.py` and adapt. +- Map core fields: layers, hidden size, FFN size, heads, KV groups, max seq len, RoPE base/scale, tied embeddings. +- Configure parallelism: `tensor_model_parallel_size`, `pipeline_model_parallel_size` (VPP/EP optional). +- Configure numerics: `fp16`/`bf16`, `params_dtype`, activation recompute. +- If needed, point to custom attention/MLP via layer spec. +- Return a lazily constructed distributed model in `.provide_distributed_model()`. + +Reference providers: +- Llama: `src/megatron/bridge/models/llama/llama_provider.py` +- Qwen: `src/megatron/bridge/models/qwen/qwen_provider.py` + +Acceptance: +- No linter errors. +- Minimal smoke test constructs a model and loads a tiny HF checkpoint via the bridge. +``` + +## 4) Define Config and Parameter Mappings + +Use the `provider_bridge` method to map Hugging Face configs to a Megatron model provider, and use `MegatronMappingRegistry` to map Megatron parameter names to Hugging Face parameter names. Start with the essentials (embeddings, final norm, QKV, MLP), then add extras (biases, rotary embeddings, experts, and vision blocks). + +- `provider_bridge`: see [model_bridge.py](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/models/conversion/model_bridge.py) +- `MegatronMappingRegistry`: see [mapping_registry.py](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/models/conversion/mapping_registry.py) +- Mapping implementations: see [param_mapping.py](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/models/conversion/param_mapping.py) +- Background: see [Bridge technical details](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/docs/bridge-tech-details.md) + +Example registration skeleton: + +```python +from megatron.core.models.gpt.gpt_model import GPTModel +from transformers import LlamaForCausalLM # replace with your HF class +from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge +from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry +from megatron.bridge.models.conversion.param_mapping import AutoMapping, QKVMapping, GatedMLPMapping +from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM +from ._provider import YourModelProvider + +@MegatronModelBridge.register_bridge(source=LlamaForCausalLM, target=GPTModel) +class YourModelBridge(MegatronModelBridge): + def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> YourModelProvider: + cfg = hf_pretrained.config + return YourModelProvider( + num_layers=cfg.num_hidden_layers, + hidden_size=cfg.hidden_size, + ffn_hidden_size=getattr(cfg, "intermediate_size", 4 * cfg.hidden_size), + num_attention_heads=cfg.num_attention_heads, + num_query_groups=getattr(cfg, "num_key_value_heads", cfg.num_attention_heads), + # set dtype flags via helper if needed + params_dtype=self.dtype_from_hf(cfg), + ... + ) + + def mapping_registry(self) -> MegatronMappingRegistry: + return MegatronMappingRegistry( + AutoMapping( + megatron_param="embedding.word_embeddings.weight", + hf_param="model.embed_tokens.weight", + ), + AutoMapping( + megatron_param="output_layer.weight", + hf_param="lm_head.weight", + ), + AutoMapping( + megatron_param="decoder.final_layernorm.weight", + hf_param="model.norm.weight", + ), + QKVMapping( + megatron_param="decoder.layers.*.self_attention.linear_qkv.weight", + q="model.layers.*.self_attn.q_proj.weight", + k="model.layers.*.self_attn.k_proj.weight", + v="model.layers.*.self_attn.v_proj.weight", + ), + GatedMLPMapping( + megatron_param="decoder.layers.*.mlp.linear_fc1.weight", + gate="model.layers.*.mlp.gate_proj.weight", + up="model.layers.*.mlp.up_proj.weight", + ), + ... + ) +``` + +Notes: +- Use `*` wildcards for per-layer patterns; the number of wildcards must match between `megatron_param` and the HF pattern(s). +- `*` typically captures layer indices; `**` can match across dots. For example, to map both `.weight` and `.bias` together: + ```python + AutoMapping( + megatron_param="output_layer.**", + hf_param="lm_head.**", + ), + ``` +- In some cases, the same module can have different Megatron parameter names depending on whether you use the Transformer Engine backend or the PyTorch backend. In that case, list both mappings, e.g., `[AutoMapping(megatron_param="te_backend_name", hf_param="hf_name"), AutoMapping(megatron_param="pytorch_backend_name", hf_param="hf_name")]`. Multiple Megatron parameters can map to the same Hugging Face parameter because, during conversion, the registry only queries the current model's module names. +- Prefer `AutoMapping` when the Megatron layer type implies the TP split automatically. +- Use `QKVMapping` for fused QKV and `GatedMLPMapping` for gate/up concatenation. + +### Suggested Cursor prompt (Bridge) [Expermental] +```text +You are working in the Megatron Bridge repo. Create `src/megatron/bridge/models//_bridge.py`. + +Goal: Implement a bridge class that connects an HF model class to a Megatron model using `MegatronModelBridge`. + +Tasks: +- Add `@MegatronModelBridge.register_bridge(source=, target=GPTModel)`. +- Implement `provider_bridge(self, hf_pretrained)` to read `hf_pretrained.config` and return `YourModelProvider(...)` with mapped fields (layers, hidden size, FFN, heads, groups, RoPE, dtype via `self.dtype_from_hf(cfg)`). +- Implement `mapping_registry(self)` returning `MegatronMappingRegistry(...)` with: + - `AutoMapping` for embeddings, final norm, output layer, 1:1 mapped weights. + - `QKVMapping` for fused QKV if applicable. + - `GatedMLPMapping` for gate/up if applicable. +- Use `*` wildcards consistently between Megatron and HF patterns. +- Add model organization to SAFE_REPOS list at `megatron.bridge.models.hf_pretrained.utils` + +References: +- `src/megatron/bridge/models/conversion/model_bridge.py` +- `src/megatron/bridge/models/conversion/mapping_registry.py` +- `src/megatron/bridge/models/conversion/param_mapping.py` +- `src/megatron/bridge/models/qwen/qwen2_bridge.py` + +Acceptance: +- HF → Megatron load completes with no missing parameters (for a tiny model). +- Megatron → HF export returns tensors with expected shapes/dtypes for several keys. +``` + +## 5) Minimal smoke test (local) + +A minimal bidirectional end-to-end check: +```python +from megatron.bridge import AutoBridge + +# HF → Megatron +bridge = AutoBridge.from_hf_pretrained("/", trust_remote_code=True) +provider = bridge.to_megatron_provider() +provider.tensor_model_parallel_size = 1 +provider.pipeline_model_parallel_size = 1 +provider.finalize() +model = provider.provide_distributed_model(wrap_with_ddp=False) +bridge.load_hf_weights(model) + +# Megatron → HF (stream a few tensors) +for i, (name, tensor) in enumerate(bridge.export_hf_weights(model, cpu=True)): + print(name, tuple(tensor.shape)) + if i > 10: + break +``` + +## 6) Validate with examples +Use the examples in `examples/conversion/` to verify bidirectional conversion and basic generation with more complex model parallel setups. + +- Generate from HF directly with the bridge +- Convert checkpoints back and forth +- Multi-GPU HF load to Megatron + +```sh +python examples/conversion/hf_to_megatron_generate_text.py --hf_model_path / --prompt "Hello" +python examples/conversion/convert_checkpoints.py import --hf-model / --megatron-path ./checkpoints/ +``` +## 7) Add tests + +Add or extend tests under `tests/functional_tests/models//` and `tests/unit_tests/models/`: + +Tests are organized in model-specific subdirectories that mirror the source structure in `src/megatron/bridge/models/`. + +- Conversion coverage: + - HF → Megatron load succeeds without missing params + - Megatron → HF export round-trips shapes and dtypes +- Provider coverage: + - Provider fields align with HF config (heads, groups, FFN size, RoPE) +- Optional numeric checks: + - Forward parity on a handful of tokens comparing HF vs Megatron outputs + +Examples to reference: +- `tests/functional_tests/models/qwen/test_qwen3_provider.py` +- `tests/functional_tests/models/qwen/test_qwen3_conversion.py` + +Run fast tests locally: +```sh +uv run pytest -q tests/functional_tests/models//test__provider.py -k your_model | cat +uv run pytest -q tests/functional_tests/models//test__conversion.py -k your_model | cat +``` + +Full suite (slower): +```sh +uv run pytest -q tests | cat +``` + +### 7.1) Model not found in CI Cache + +Megatron Bridge functional tests run with `HF_HUB_OFFLINE=1`. This means that contributions including a new bridge and tests +for a HuggingFace model that is not cached in our CI's `$HF_HOME` directory will fail with an error similar to: + +``` +huggingface_hub.errors.LocalEntryNotFoundError: Cannot find the requested files in the disk cache and outgoing traffic has been disabled. +``` + +If such an error is encountered in the CI, please request a repo maintainer to launch the 'Cache HuggingFace model' workflow for the model(s) +you are adding support for in your PR. + +### Suggested Cursor prompt (Tests) [Expermental] +```text +You are working in the Megatron Bridge repo. Add tests for a new model ``. + +Create a subdirectory `tests/functional_tests/models//` with an `__init__.py` file and two test modules: +1) `test__provider.py` + - Build a tiny HF model/config (or use `/` if available). + - Use the bridge to derive a provider and construct the model with TP=PP=1. + - Assert provider fields match HF config (heads, groups, hidden size, FFN, RoPE, vocab size, max position). + +2) `test__conversion.py` + - HF → Megatron: load HF weights into the Megatron model via the bridge; assert no missing/extra params. + - Megatron → HF: export a subset of tensors; assert shape/dtype parity with HF. + - Optionally run a short generation on CPU and compare logits numerically within tolerance. + +Use `tests/functional_tests/models/qwen/test_qwen3_provider.py` and `test_qwen3_conversion.py` as templates. + +Provide `-k your_model` selectors and guard long tests with `pytest.skip` if external weights are unavailable. +``` + +## 8) Troubleshooting + +- Shape mismatches: double-check TP/PP splits and model configs. +- Missing weights: ensure every Megatron param has a mapping; print unresolved names. +- Dtype issues: cast HF weights to destination dtype inside mappings when needed. +- EP/MoE layers: see EP-specific gather/scatter helpers in `param_mapping.py`. + +Enable verbose logs: +```python +import logging +logging.getLogger("megatron.bridge").setLevel(logging.DEBUG) +``` + +## 9) PR checklist + +- Provde details in PR descriptions +- Provider maps all required config fields +- All parameters are covered by mappings +- Generation results after conversion from HF to Megatron match Megatron, including multi-GPU runs +- Unit/functional tests added and green +- Add your model to the Supported Models table in the repo `README.md` if applicable + +## 10) Useful links + +- User guide: [docs/bridge-guide.md](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/docs/bridge-guide.md) +- Technical deep-dive: [docs/bridge-tech-details.md](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/docs/bridge-tech-details.md) +- Code examples: [examples/conversion/](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/examples/conversion) +- Providers and bridges: [src/megatron/bridge/models/](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models) +- GitHub source tree: [Megatron Bridge src/megatron/bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge) diff --git a/fern/v0.2.0/pages/bridge-guide.mdx b/fern/v0.2.0/pages/bridge-guide.mdx new file mode 100644 index 0000000000..e934ddbef3 --- /dev/null +++ b/fern/v0.2.0/pages/bridge-guide.mdx @@ -0,0 +1,315 @@ +--- +title: "Get Started with 🤗 Hugging Face Conversion" +description: "" +--- + +Megatron Bridge provides seamless bidirectional conversion between 🤗 Hugging Face Transformers and Megatron model definitions. This guide covers the main APIs for loading models, checking compatibility, and converting between formats. + +## Design and Goals + +- Single high-level entry point: `AutoBridge` detects HF model architectures and dispatches to the correct bridge. +- Bidirectional conversion: Import HF → Megatron for training; export Megatron → HF for deployment. +- Parallelism-aware: Handles TP/PP/VPP/CP/EP/ETP distributions during conversion. +- Streaming and memory efficiency: per-parameter streaming using safetensors. +- Provider pattern: Configure Megatron-Core `TransformerConfig`-compatible attributes before instantiation via `to_megatron_provider()`. +- Convenience workflows: `import_ckpt` and `export_ckpt` provide one-call HF↔Megatron checkpoint flows. + +See the repository `README.md` for installation, supported models, and project highlights. + +## Loading a 🤗 Hugging Face Model into Megatron + +The easiest way to load a 🤗 Hugging Face model is using `AutoBridge.from_hf_pretrained()`, which automatically detects the model architecture and selects the appropriate bridge for conversion. You can then use `AutoBridge.to_megatron_model()` to initialize the Megatron model from the 🤗 Hugging Face configuration and load 🤗 Hugging Face weights at the same time. + +### Accessing Gated 🤗 Hugging Face Models + +Some models in Megatron Bridge require access to gated repositories on Hugging Face. These are models that require explicit permission from the model authors before you can download or use them. + +If you encounter an error like this when trying to use a model: + +``` +OSError: You are trying to access a gated repo. +Make sure to have access to it at +``` + +Follow these steps to resolve the issue: + +1. **Request access**: Visit the URL provided in the error message and request access to the gated model +2. **Generate a token**: Create a Hugging Face access token by following [this tutorial](https://huggingface.co/docs/hub/en/security-tokens#how-to-manage-user-access-tokens) +3. **Set the environment variable**: Export your token in your environment: + +```bash +export HF_TOKEN= +``` + +### Basic Usage + +```python +from megatron.bridge import AutoBridge + +# Load a supported model automatically +bridge = AutoBridge.from_hf_pretrained("meta-llama/Llama-3.2-1B") + +# Create a provider, configure before instantiation, then build the model +provider = bridge.to_megatron_provider() +provider.tensor_model_parallel_size = 1 +provider.pipeline_model_parallel_size = 1 +provider.finalize() +megatron_model = provider.provide_distributed_model(wrap_with_ddp=False) +``` + +### Advanced Loading Options +You can also load models with specific settings such as precision, device placement, or by enabling trust in remote code: + +```python +import torch +from megatron.bridge import AutoBridge + +# Load with specific settings +bridge = AutoBridge.from_hf_pretrained( + "meta-llama/Llama-2-7b-hf", + torch_dtype=torch.float16, + device_map="auto", + trust_remote_code=True, +) + +# Load from local path +bridge = AutoBridge.from_hf_pretrained("/path/to/local/hf_model") +``` + +### Using Model Providers + +For more control over model configuration, use the provider pattern. The provider lets you configure any `TransformerConfig` attribute: + +```python +from megatron.bridge import AutoBridge + +# Load a supported model automatically +bridge = AutoBridge.from_hf_pretrained("meta-llama/Llama-3.2-1B") + +# Get a model provider (lazy loading) +provider = bridge.to_megatron_provider() + +# Configure parallelism (multi-GPU requires torchrun or srun) +provider.tensor_model_parallel_size = 8 +provider.pipeline_model_parallel_size = 2 + +# Configure fusions +provider.bias_activation_fusion = True +provider.bias_dropout_fusion = True + +# Finalize the provider to run validation checks and complete initialization +provider.finalize() + +# Create the model with all configurations applied +model = provider.provide_distributed_model(wrap_with_ddp=False) +``` + +The provider pattern is especially useful when you need to: +- Override default model parameters +- Configure advanced features like MoE, activation recomputation, or mixed precision +- Set up distributed training parameters + +## Check Supported Models + +Before loading a model, you can check if it's supported by Megatron Bridge. + +You can list all supported 🤗 Hugging Face model architectures as follows: + +```python +from megatron.bridge import AutoBridge + +# Get a list of all supported model architectures +supported_models = AutoBridge.list_supported_models() + +print(f"Found {len(supported_models)} supported models:") +for i, model in enumerate(supported_models, 1): + print(f" {i:2d}. {model}") +``` + +Alternatively, check if a specific model is supported: + +```python +from megatron.bridge import AutoBridge + +if AutoBridge.can_handle("meta-llama/Llama-3.2-1B"): + print("✅ Model is supported!") + bridge = AutoBridge.from_hf_pretrained("meta-llama/Llama-3.2-1B") +else: + print("❌ Model requires a custom bridge implementation") +``` + +## Converting Back to 🤗 Hugging Face + +After training or modifying a Megatron model, you can convert it back to 🤗Hugging Face format for deployment or sharing. The bridge provides several methods for this conversion depending on your needs. + +To save the complete model including configuration, tokenizer, and weights: + +```python +# Save the complete model (config, tokenizer, weights) +bridge.save_hf_pretrained(megatron_model, "./my-fine-tuned-llama") + +# The saved model can be loaded with 🤗 Hugging Face +from transformers import AutoModelForCausalLM +hf_model = AutoModelForCausalLM.from_pretrained("./my-fine-tuned-llama") +``` + +You can save the model weights (safetensors): + +```python +# Save just the model weights (faster, smaller) +bridge.save_hf_weights(megatron_model, "./model_weights") + +# Save without progress bar (useful in scripts) +bridge.save_hf_weights(megatron_model, "./weights", show_progress=False) +``` + +You can also stream weights without saving to disk during conversion for on-the-fly use in RL frameworks, for example: + +```python +# Stream weights during conversion (memory efficient) +for name, weight in bridge.export_hf_weights(megatron_model): + print(f"Exporting {name}: {weight.shape}") + +for name, weight in bridge.export_hf_weights(megatron_model, cpu=True): + print(f"Exported {name}: {tuple(weight.shape)}") +``` + +## Common Patterns and Best Practices +When working with Megatron Bridge, there are several patterns that will help you use the API effectively and avoid common pitfalls. + +### 1. Always Use High-Level APIs +Always prefer high-level APIs like `AutoBridge` for automatic model detection. Avoid direct bridge usage unless you know the specific type required: + +```python +# ✅ Preferred: Use AutoBridge for automatic detection +bridge = AutoBridge.from_hf_pretrained("any-supported-model") + +# ❌ Avoid: Direct bridge usage unless you know the specific type +``` + +### 2. Configure Before Creating Models +When using the provider pattern, always configure parallelism and other settings before creating the model. Creating the model first uses default settings that may not be optimal: + +```python +# ✅ Correct: Configure provider before creating model +provider = bridge.to_megatron_provider() +provider.tensor_model_parallel_size = 8 +provider.finalize() +model = provider.provide_distributed_model(wrap_with_ddp=False) + +# ❌ Avoid: Creating model before configuring parallelism +model = bridge.to_megatron_model() # Uses default settings +``` + +### 3. Leverage the Parameter Streaming API +You can stream converted weights from Megatron to HF without saving to disk: + +```python +# ✅ Use streaming for large models +for name, weight in bridge.export_hf_weights(model, cpu=True): + process_weight(name, weight) +``` + +### 4. Use `from_hf_pretrained` for Export Workflows + +When exporting Megatron checkpoints back to 🤗 Hugging Face format, always use `from_hf_pretrained()` instead of `from_hf_config()`. The `from_hf_config()` method does not load the tokenizer and other artifacts required for saving a complete 🤗 Hugging Face checkpoint: + +```python +from megatron.bridge import AutoBridge + +# ✅ Correct: Use from_hf_pretrained for export workflows +bridge = AutoBridge.from_hf_pretrained("meta-llama/Llama-3.2-1B") +bridge.export_ckpt("./megatron_checkpoints/llama32_1b", "./hf_exports/llama32_1b") + +# ❌ Avoid: from_hf_config lacks artifacts needed for saving +# config = AutoConfig.from_pretrained("meta-llama/Llama-3.2-1B") +# bridge = AutoBridge.from_hf_config(config) # Missing tokenizer, etc. +# bridge.export_ckpt(...) # Will fail! +``` + +The `from_hf_config()` method is only suitable for architecture exploration and introspection (e.g., inspecting `transformer_config`), not for checkpoint conversion workflows. + +For more examples and advanced usage patterns, see the `examples/conversion/` directory in the repository. + +## Convenience Workflows (Commands) + +These examples can be run directly as shell commands. + +### HF → Megatron checkpoint import (one call) + +```bash +huggingface-cli login --token +python -c "from megatron.bridge import AutoBridge; AutoBridge.import_ckpt('meta-llama/Llama-3.2-1B','./megatron_checkpoints/llama32_1b')" +``` + +### Megatron → HF export (one call) + +```bash +python -c "from megatron.bridge import AutoBridge; b=AutoBridge.from_hf_pretrained('meta-llama/Llama-3.2-1B'); b.export_ckpt('./megatron_checkpoints/llama32_1b','./hf_exports/llama32_1b')" +``` + +### Create Megatron models and run locally + +```bash +python - << 'PY' +from megatron.bridge import AutoBridge + +bridge = AutoBridge.from_hf_pretrained('meta-llama/Llama-3.2-1B') +provider = bridge.to_megatron_provider() +provider.tensor_model_parallel_size = 1 +provider.pipeline_model_parallel_size = 1 +provider.finalize() +model = provider.provide_distributed_model(wrap_with_ddp=False) + +# Export to HF folder +bridge.save_hf_pretrained(model, './hf_exports/llama32_1b') +PY +``` + +### Launch with multiple GPUs (example) + +```bash +torchrun --nproc-per-node=2 -m examples.conversion.generate_from_hf +``` + +## AutoBridge API Reference + +Latest public APIs and signatures (see [API Documentation](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/)): + +```python +from megatron.bridge import AutoBridge + +# Creation and capability +AutoBridge.from_hf_pretrained(path: str | Path, **kwargs) -> AutoBridge +AutoBridge.from_hf_config(config: PretrainedConfig) -> AutoBridge +AutoBridge.can_handle(path: str | Path, trust_remote_code: bool = False) -> bool +AutoBridge.list_supported_models() -> list[str] +AutoBridge.supports(config: Any) -> bool + +# Provider/model construction +AutoBridge.to_megatron_provider(load_weights: bool = True, hf_path: str | Path | None = None) -> GPTModelProvider +AutoBridge.to_megatron_model(load_weights: bool = True, hf_path: str | Path | None = None, **kwargs) -> list[MegatronModule] + +# HF → Megatron weights +AutoBridge.load_hf_weights(model: list[MegatronModule], hf_path: str | Path | None = None) -> None + +# Megatron → HF conversion +AutoBridge.export_hf_weights(model: list[MegatronModule], cpu: bool = False, show_progress: bool = True, conversion_tasks: Optional[list[WeightConversionTask]] = None) -> Iterable[HFWeightTuple] +AutoBridge.save_hf_pretrained(model: list[MegatronModule], path: str | Path, show_progress: bool = True) -> None +AutoBridge.save_hf_weights(model: list[MegatronModule], path: str | Path, show_progress: bool = True) -> None + +# Megatron native checkpoints +AutoBridge.save_megatron_model(model: list[MegatronModule], path: str | Path) -> None +AutoBridge.load_megatron_model(path: str | Path, **kwargs) -> list[MegatronModule] + +# One-call workflows +AutoBridge.import_ckpt(hf_model_id: str | Path, megatron_path: str | Path, **kwargs) -> None # HF → Megatron ckpt +AutoBridge.export_ckpt(megatron_path: str | Path, hf_path: str | Path, show_progress: bool = True) -> None # Megatron → HF + +# Config extraction +AutoBridge.transformer_config -> TransformerConfig +AutoBridge.mla_transformer_config -> MLATransformerConfig + +# Introspection / planning +AutoBridge.get_conversion_tasks(megatron_model: MegatronModule | list[MegatronModule], hf_path: str | Path | None = None) -> list[WeightConversionTask] +``` diff --git a/fern/v0.2.0/pages/bridge-rl-integration.mdx b/fern/v0.2.0/pages/bridge-rl-integration.mdx new file mode 100644 index 0000000000..db7943955f --- /dev/null +++ b/fern/v0.2.0/pages/bridge-rl-integration.mdx @@ -0,0 +1,498 @@ +--- +title: "Adapting Megatron Bridge in Reinforcement Learning Frameworks" +description: "" +--- + +Megatron Bridge provides a clean, parallelism-aware path to use 🤗 Hugging Face models with Megatron-Core training and convert back again for inference. This guide shows how to adapt Megatron Bridge into a new RL framework to: + +- Convert Hugging Face (HF) checkpoints → Megatron format for scalable training +- Train with Megatron-Core using TP/PP/CP/MoE parallelism, checkpointing, and efficient data paths +- Refit trained weights back to HF for deployment with inference engines (e.g., vLLM), including zero-copy/IPC flows + +The examples mirror how NeMo-RL integrates Megatron Bridge: + +- [nemo_rl/models/megatron/community_import.py](https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/models/megatron/community_import.py) +- [nemo_rl/models/policy/megatron_policy_worker.py](https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/models/policy/megatron_policy_worker.py) + +- Local example script in this repo: [examples/rl/rlhf_with_bridge.py](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/rl/rlhf_with_bridge.py) + +## Prerequisites + +- A working PyTorch + NCCL GPU stack +- Megatron-LM (MCore) and Megatron-Bridge installed +- A distributed launcher (e.g., `torchrun`, `srun`) for multi-GPU setups +- Access to gated HF repos if needed (export `HF_TOKEN`) + +```bash +export HF_TOKEN= +``` + +## 1) One-time HF → Megatron checkpoint conversion + +Use `AutoBridge` to import an HF model into Megatron format. This writes a Megatron checkpoint directory with a `run_config.yaml` you will reuse during training. + +```python +from megatron.bridge import AutoBridge + +# Import a model to Megatron checkpoint format (one call) +AutoBridge.import_ckpt( + hf_model_id="meta-llama/Llama-3.2-1B", + megatron_path="/path/to/megatron_ckpt/llama32_1b", +) +``` + +Or, with explicit provider and parallelism settings (similar to [nemo_rl/models/megatron/community_import.py](https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/models/megatron/community_import.py)): + +```python +from megatron.bridge import AutoBridge + +bridge = AutoBridge.from_hf_pretrained("meta-llama/Llama-3.2-1B", trust_remote_code=True) +provider = bridge.to_megatron_provider(load_weights=True) + +# Configure distributed parallelism used during IMPORT +provider.tensor_model_parallel_size = 2 +provider.pipeline_model_parallel_size = 1 +provider.expert_model_parallel_size = 1 +provider.expert_tensor_parallel_size = 1 +provider.num_layers_in_first_pipeline_stage = 0 +provider.num_layers_in_last_pipeline_stage = 0 +provider.finalize() + +# Create distributed model and save as Megatron checkpoint +megatron_model = provider.provide_distributed_model(wrap_with_ddp=False) +bridge.save_megatron_model(megatron_model, "/path/to/megatron_ckpt") +``` + +You can also check and try out our multi-GPU conversion example script: [examples/conversion/hf_megatron_roundtrip_multi_gpu.py](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/conversion/hf_megatron_roundtrip_multi_gpu.py) + +Notes: +- The import-time parallelism is only for loading/conversion. The saved config is restored to canonical values to avoid validation issues at training time. +- If you are running inside a framework, make sure to clean up any existing distributed state before and after import by destroying or initializing process groups as needed. The `provide_distributed_model` method will initialize a new distributed environment if one is not already set up. + +## 2) Build training configuration and initialize Megatron-Core + +Translate your RL framework config into Megatron Bridge's `ConfigContainer` for model, optimizer, scheduler, DDP, tokenizer, and checkpoints. + +```python +import torch +from megatron.bridge.training.config import ( + ConfigContainer, + TrainingConfig, + OptimizerConfig, + SchedulerConfig, + DistributedDataParallelConfig, + CheckpointConfig, + TokenizerConfig, +) +from nemo_rl.models.policy import PolicyConfig # or your own policy cfg type + +# Example: map your RL config to Megatron config +def build_megatron_config(rl_cfg: PolicyConfig, pretrained_ckpt_dir: str) -> ConfigContainer: + model_cfg = rl_cfg["megatron_cfg"].copy() + # Precision + dtype = rl_cfg["precision"] + model_cfg["bf16"] = dtype == "bfloat16" + model_cfg["fp16"] = dtype == "float16" + + checkpoint = CheckpointConfig( + save_interval=100, + save=rl_cfg["train_ckpt_dir"], + load=rl_cfg["train_ckpt_dir"], + pretrained_checkpoint=pretrained_ckpt_dir, + async_save=False, + fully_parallel_save=True, + fully_parallel_load=True, + load_rng=False, + ) + + ddp = DistributedDataParallelConfig( + check_for_nan_in_grad=True, + grad_reduce_in_fp32=rl_cfg["megatron_cfg"]["distributed_data_parallel_config"]["grad_reduce_in_fp32"], + overlap_grad_reduce=rl_cfg["megatron_cfg"]["distributed_data_parallel_config"]["overlap_grad_reduce"], + overlap_param_gather=rl_cfg["megatron_cfg"]["distributed_data_parallel_config"]["overlap_param_gather"], + average_in_collective=rl_cfg["megatron_cfg"]["distributed_data_parallel_config"]["average_in_collective"], + use_distributed_optimizer=rl_cfg["megatron_cfg"]["optimizer"]["use_distributed_optimizer"], + data_parallel_sharding_strategy=rl_cfg["megatron_cfg"]["distributed_data_parallel_config"]["data_parallel_sharding_strategy"], + ) + + opt = OptimizerConfig(**rl_cfg["megatron_cfg"]["optimizer"]) # lr, wd, etc. + sch = SchedulerConfig(**rl_cfg["megatron_cfg"]["scheduler"]) # warmup, decay, etc. + + train = TrainingConfig( + micro_batch_size=rl_cfg["train_micro_batch_size"], + global_batch_size=rl_cfg["train_global_batch_size"], + train_iters=rl_cfg["megatron_cfg"]["train_iters"], + ) + + tokenizer = TokenizerConfig( + tokenizer_type="HuggingFaceTokenizer", + tokenizer_model=rl_cfg["model_name"], + ) + + return ConfigContainer( + model=model_cfg, + checkpoint=checkpoint, + logger=None, + train=train, + optimizer=opt, + ddp=ddp, + scheduler=sch, + dataset=None, + tokenizer=tokenizer, + ) +``` + +Initialize Megatron-Core using a helper similar to `setup_megatron_model` from NeMo-RL: + +```python +from megatron.bridge.training.initialize import initialize_megatron, set_jit_fusion_options +from megatron.bridge.models.model_provider import get_model +from megatron.bridge.training.optim import setup_optimizer +from megatron.bridge.training.checkpointing import init_checkpointing_context, load_checkpoint +from megatron.bridge.training.state import GlobalState + +# Minimal bootstrap +state = GlobalState() +state.cfg = megatron_cfg +initialize_megatron(cfg=megatron_cfg) +set_jit_fusion_options(megatron_cfg.model, megatron_cfg.train.micro_batch_size) + +ckpt_ctx = init_checkpointing_context(megatron_cfg.checkpoint) +model_list = get_model( + megatron_cfg.model, + megatron_cfg.ddp, + use_torch_fsdp2=megatron_cfg.dist.use_torch_fsdp2, + overlap_param_gather_with_optimizer_step=megatron_cfg.optimizer.overlap_param_gather_with_optimizer_step, + data_parallel_random_init=megatron_cfg.rng.data_parallel_random_init, +) +optimizer, scheduler = setup_optimizer( + optimizer_config=megatron_cfg.optimizer, + scheduler_config=megatron_cfg.scheduler, + model=model_list, + use_gloo_process_groups=megatron_cfg.dist.use_gloo_process_groups, +) + +# Optional: load pretrained checkpoint +load_checkpoint( + state, + model_list, + optimizer, + scheduler, + checkpointing_context=ckpt_ctx, + skip_load_to_model_and_opt=False, +) + +model = model_list[0] +``` + +Key mappings to handle in your RL config: +- Parallelism: `tensor_model_parallel_size`, `pipeline_model_parallel_size`, `context_parallel_size` (requires sequence packing), and MoE (`expert_*`). +- Precision: `bf16`/`fp16` plus `pipeline_dtype`. +- Activation checkpointing: recompute settings for memory savings. +- FP8 (advanced): be mindful of alignment/padding requirements if enabled. + +## 3) Training loop integration (forward/backward/microbatches) + +Megatron-Core exposes `get_forward_backward_func()` to run a microbatch loop. Plug in your RL loss function. + +```python +from functools import partial +from megatron.core.pipeline_parallel import get_forward_backward_func + +model.train() +forward_backward = get_forward_backward_func() + +# Your loss function should return (loss_tensor, metrics_dict) +def rl_loss_fn(outputs, batch): + # Compute logits → loss for your RL objective (e.g., PPO, DPO) + loss = outputs.sum() * 0.0 # placeholder + return loss, {"loss": loss.detach()} + +# Forward step: prepare inputs; return outputs and a collector that yields loss + +def forward_step_fn(data_iterator, model): + batch = next(data_iterator).to("cuda") + outputs = model( + input_ids=batch["input_ids"], + attention_mask=batch.get("attention_mask"), + packed_seq_params=batch.get("packed_seq_params"), # if sequence packing + # multimodal features can be passed as kwargs + ) + return outputs, (lambda _out: rl_loss_fn(outputs, batch)) + +losses_reduced = forward_backward( + forward_step_func=forward_step_fn, + data_iterator=your_microbatch_iterator, + model=model, + num_microbatches=num_microbatches, + seq_length=sequence_length, + micro_batch_size=micro_batch_size, + decoder_seq_length=sequence_length, + forward_only=False, + do_not_average_loss=True, +) + +# Optimizer/scheduler steps +update_successful, grad_norm, _ = optimizer.step() +scheduler.step(increment=global_batch_size) +``` + +Sequence packing and context parallelism: +- If `context_parallel_size > 1`, enable sequence packing and build `packed_seq_params` and `cu_seqlens` per microbatch before calling the model. +- With FP8, ensure sequence padding respects hardware-friendly multiples (e.g., lcm(16, 2 × TP × CP)). + +## 4) Token logprobs for RL objectives (advantages, DPO, etc.) + +For evaluation of token logprobs, run forward-only and reduce TP-sharded logits to per-token logprobs. + +```python +import torch +from megatron.core.parallel_state import get_tensor_model_parallel_group, get_tensor_model_parallel_rank + +@torch.no_grad() +def get_token_logprobs(model, batch): + model.eval() + input_ids = batch["input_ids"].to("cuda") + outputs = model(input_ids=input_ids) + + # Reduce TP logits → local logprobs for targets + tp_group = get_tensor_model_parallel_group() + tp_rank = get_tensor_model_parallel_rank() + + # Use a reducer similar to NeMo-RL's `from_parallel_logits_to_logprobs` + token_logprobs = your_reduce_parallel_logits_to_logprobs( + outputs, + target=input_ids, + vocab_start_index=tp_rank * outputs.shape[-1], + vocab_end_index=(tp_rank + 1) * outputs.shape[-1], + tp_group=tp_group, + inference_only=True, + ) + + # Prepend a zero to keep the same sequence length as the inputs + token_logprobs = torch.cat([torch.zeros_like(token_logprobs[:, :1]), token_logprobs], dim=1) + return token_logprobs.cpu() +``` + +If using sequence packing + context parallelism, switch to the packed variant that uses `packed_seq_params` and `cu_seqlens` for correct alignment. + +## 5) Checkpointing (save/load) + +Use Megatron-Bridge’s checkpoint helpers. Temporarily disable overlapping param-gather hooks if needed during save. + +```python +from megatron.bridge.training.checkpointing import ( + save_checkpoint, + load_checkpoint, + init_checkpointing_context, +) + +ckpt_ctx = init_checkpointing_context(megatron_cfg.checkpoint) +save_checkpoint( + state=state, + model=[model], + optimizer=optimizer, + opt_param_scheduler=scheduler, + num_floating_point_operations_so_far=state.train_state.floating_point_operations_so_far, + checkpointing_context=ckpt_ctx, +) +``` + +Tips: +- Prefer fully-parallel save/load at scale (`fully_parallel_save=True`, `fully_parallel_load=True`). + +## 6) Refit: Megatron → HF for inference (vLLM, Triton, etc.) + +Two common pathways: + +### A) Export full HF checkpoint (simplest) + +```python +from megatron.bridge import AutoBridge + +bridge = AutoBridge.from_hf_pretrained("meta-llama/Llama-3.2-1B", trust_remote_code=True) +# Load Megatron model from your training checkpoint +megatron_model = bridge.load_megatron_model("/path/to/train_ckpt") + +# Iterate over HF weights parameter-by-parameter +for name, weight in bridge.export_hf_weights(megatron_model, cpu=True, show_progress=False): + # process_or_save(name, weight) + pass +``` + +Point your inference engine (e.g., vLLM) to `"/path/to/hf_export"`. + +### B) Zero-copy streaming via ZMQ (fast refit, colocated) + +Stream tensors from the training side to your inference runtime without writing to disk. The transport is ZMQ peer-to-peer with async send/recv and ping‑pong buffers for overlap; Ray is used only for lightweight coordination. This replaces the earlier ad‑hoc per‑tensor IPC handle passing and aligns with the refactor in [NVIDIA-NeMo/RL#1267](https://github.com/NVIDIA-NeMo/RL/pull/1267). + +**Concepts (how the plan and chunking work):** +- **Transport and overlap:** ZMQ P2P streaming with asynchronous send/recv and ping‑pong buffers enables overlap between gathering and applying weights. +- **Conversion tasks (planning):** `bridge.get_conversion_tasks([model])` returns an ordered list of per-parameter conversion tasks that encode how to transform sharded Megatron weights (TP/PP/MoE/CP) back to HF tensors. The worker stores this in `self.refit_conversion_tasks` and advances a cursor (`self.refit_conversion_tasks_current_index`) as chunks are streamed. See `nemo_rl/models/policy/megatron_policy_worker.py` methods `prepare_refit_info()`, `_calculate_refit_param_info()`, and `get_weights_ipc_handles()`. +- **Size estimation across PP ranks:** Parameters are only materialized on their owning PP rank. The worker computes per-parameter byte sizes and then broadcasts those sizes to all PP ranks so the entire pipeline can agree on chunk boundaries. See `broadcast_object_across_pp_ranks()` and `_calculate_refit_param_info()` in `megatron_policy_worker.py`. +- **Memory-aware chunking:** Use your free GPU memory budget (e.g., `NRL_REFIT_BUFFER_MEMORY_RATIO`) to decide how many parameters to include in the next chunk (the set of `keys`). The worker exposes `prepare_weights_for_ipc()` which returns `(param_info, total_available_bytes)` and resets the conversion cursor; then the controller repeatedly selects `keys` whose cumulative byte size ≤ budget and streams them to the consumer over ZMQ. +- **Device routing:** Handles are returned under a `device_uuid` key (NVML UUID of the CUDA device). The inference side should map handles on the same device (or coordinate via your communicator). For collective updates, the worker can also broadcast tensors directly (`broadcast_weights_for_collective`). +- **Parallelism nuances:** With TP/EP, exported HF tensors are reassembled from shards; with CP/sequence packing, shapes/dtypes are already consistent at export time. FP8 or mixed precision can affect size estimates; the worker accounts for dtype scaling when estimating bytes. + +```python +import os +import torch +from collections import defaultdict +from megatron.bridge import AutoBridge + +bridge = AutoBridge.from_hf_pretrained("meta-llama/Llama-3.2-1B", trust_remote_code=True) + +# 1) Plan: inspect names/shapes/dtypes and estimate memory +refit_param_info_hf = {} +for name, tensor in bridge.export_hf_weights([model], show_progress=False): + refit_param_info_hf[name] = (tuple(tensor.shape), tensor.dtype) + +# 2) Budget for staging buffers (optionally) +from nemo_rl.utils.nvml import get_free_memory_bytes # or your own NVML wrapper +free_bytes = get_free_memory_bytes(torch.cuda.current_device()) +ratio = float(os.getenv("NRL_REFIT_BUFFER_MEMORY_RATIO", "0.2")) +allowed_bytes = int(free_bytes * ratio) + +# 3) Stream chunks over ZMQ +from nemo_rl.utils.nvml import get_device_uuid + +# Build conversion tasks once and advance an index as you stream +refit_conversion_tasks = bridge.get_conversion_tasks([model]) +refit_tasks_current_index = 0 + +def stream_next_chunk(keys: list[str]): + """Yield ZMQ multipart frames for this chunk. + Frames typically include: (metadata_json_bytes, payload_bytes). + """ + global refit_tasks_current_index + conversion_tasks = refit_conversion_tasks[ + refit_tasks_current_index : refit_tasks_current_index + len(keys) + ] + refit_tasks_current_index += len(keys) + + device_uuid = get_device_uuid(torch.cuda.current_device()) + + # Worker exposes a streaming generator that overlaps gather and send + for frames in worker.stream_refit_chunks( + conversion_tasks=conversion_tasks, device_uuid=device_uuid + ): + yield frames # send via zmq_socket.send_multipart(frames) + +# Example usage (producer) +for frames in stream_next_chunk(list(refit_param_info_hf.keys())): + zmq_socket.send_multipart(frames) +``` + +**Chunking in practice (controller-side selection of keys):** + +```python +# param_info like [(name, size_bytes), ...] from prepare_refit_info or prepare_weights_for_ipc +param_info, budget_bytes = worker.prepare_weights_for_ipc() + +cursor = 0 +while cursor < len(param_info): + batch_keys = [] + used = 0 + # Greedy pack parameters into this chunk until we run out of budget + while cursor < len(param_info): + name, size_bytes = param_info[cursor] + # size_bytes is already broadcast to all PP ranks; can be int + if used + int(size_bytes) > budget_bytes and len(batch_keys) > 0: + break + batch_keys.append(name) + used += int(size_bytes) + cursor += 1 + + # Stream this chunk and consume on the inference side + for frames in worker.stream_refit_chunks(keys=batch_keys): + zmq_socket.send_multipart(frames) +``` + +Environment knobs: +- `NRL_REFIT_BUFFER_MEMORY_RATIO` (default `0.2`) — fraction of free GPU memory to plan staging + +## 7) Minimal adapter skeleton + +Use this skeleton to embed Megatron Bridge into your RL codebase. Fill in the config mapping, microbatching, and loss logic. + +```python +import torch +from megatron.bridge.training.config import ConfigContainer +from megatron.bridge.training.state import GlobalState +from megatron.bridge.training.checkpointing import init_checkpointing_context, save_checkpoint +from megatron.core.pipeline_parallel import get_forward_backward_func + +class MegatronBridgeAdapter: + def __init__(self, rl_cfg, pretrained_ckpt_dir: str): + self.rl_cfg = rl_cfg + self.megatron_cfg = build_megatron_config(rl_cfg, pretrained_ckpt_dir) + self.state = GlobalState(); self.state.cfg = self.megatron_cfg + self.ckpt_ctx = init_checkpointing_context(self.megatron_cfg.checkpoint) + self._init_model() + + def _init_model(self): + from megatron.bridge.training.initialize import initialize_megatron, set_jit_fusion_options + from megatron.bridge.models.model_provider import get_model + from megatron.bridge.training.optim import setup_optimizer + initialize_megatron(cfg=self.megatron_cfg) + set_jit_fusion_options(self.megatron_cfg.model, self.megatron_cfg.train.micro_batch_size) + self.model_list = get_model(self.megatron_cfg.model, self.megatron_cfg.ddp, + use_torch_fsdp2=self.megatron_cfg.dist.use_torch_fsdp2, + overlap_param_gather_with_optimizer_step=self.megatron_cfg.optimizer.overlap_param_gather_with_optimizer_step) + self.model = self.model_list[0] + self.optimizer, self.scheduler = setup_optimizer(self.megatron_cfg.optimizer, self.megatron_cfg.scheduler, self.model_list, + use_gloo_process_groups=self.megatron_cfg.dist.use_gloo_process_groups) + + @torch.no_grad() + def get_logprobs(self, batch): + self.model.eval() + # Implement reduction from parallel logits to token logprobs + ... + + def train_step(self, mb_iter, num_microbatches, seq_len, mbs, loss_fn): + self.model.train() + fb = get_forward_backward_func() + def fwd(data_it, model): + batch = next(data_it).to("cuda") + out = model(input_ids=batch["input_ids"], attention_mask=batch.get("attention_mask")) + return out, (lambda _o: loss_fn(out, batch)) + fb(forward_step_func=fwd, data_iterator=mb_iter, model=self.model, num_microbatches=num_microbatches, + seq_length=seq_len, micro_batch_size=mbs, decoder_seq_length=seq_len, forward_only=False, do_not_average_loss=True) + ok, _, _ = self.optimizer.step(); self.scheduler.step(increment=self.rl_cfg["train_global_batch_size"]) + return ok + + def save_ckpt(self, path: str): + save_checkpoint(self.state, [self.model], self.optimizer, self.scheduler, + num_floating_point_operations_so_far=self.state.train_state.floating_point_operations_so_far, + checkpointing_context=self.ckpt_ctx) + + def export_hf(self, out_dir: str, trust_remote_code: bool = False): + from megatron.bridge import AutoBridge + bridge = AutoBridge.from_hf_pretrained(self.rl_cfg["model_name"], trust_remote_code=trust_remote_code) + # Stream weights directly using AutoBridge.export_hf_weights; consume (save/IPC) as needed + for name, tensor in bridge.export_hf_weights([self.model], show_progress=False): + # process_or_save(name, tensor, out_dir) # implement your consumer (e.g., safetensors or IPC) + pass + # Optionally, to persist safetensors shards without config/tokenizer: + # bridge.save_hf_weights([self.model], out_dir, show_progress=False) +``` + +## 8) Best practices and pitfalls + +- Parallelism + - If `context_parallel_size > 1`, enable sequence packing. Use packed logprob reducers. + - With FP8, pad to hardware-friendly multiples (e.g., lcm(16, 2 × TP × CP)). +- Offloading + - TBA +- MoE router stability + - Consider freezing router weights and disabling router load balancing to reduce training instability (see `freeze_moe_router`, `moe_router_bias_update_rate=0.0`). +- Hooks + - Temporarily disable DDP forward pre-hooks when swapping weights or saving to avoid conflicts with overlapping param gather. +- Checkpointing + - Use fully-parallel save/load at scale. Avoid async save unless validated in your environment. + +## See also + +- [Bridge with 🤗 Hugging Face](/./bridge-guide) for HF↔Megatron conversion overview +- [nemo_rl/models/megatron/community_import.py](https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/models/megatron/community_import.py) for import/export helpers +- [nemo_rl/models/policy/megatron_policy_worker.py](https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/models/policy/megatron_policy_worker.py) for end-to-end RL integration (training, logprobs, generation, refit) diff --git a/fern/v0.2.0/pages/bridge-tech-details.mdx b/fern/v0.2.0/pages/bridge-tech-details.mdx new file mode 100644 index 0000000000..cd93b8f963 --- /dev/null +++ b/fern/v0.2.0/pages/bridge-tech-details.mdx @@ -0,0 +1,299 @@ +--- +title: "Megatron Bridge Conversion Technical Details" +description: "" +--- + +Megatron Bridge provides a robust, parallelism-aware pathway to convert models and checkpoints between 🤗 Hugging Face Transformers and Megatron-Core formats. This page dives into the architecture, data flows, and per-parameter conversion engine, with examples. + +Megatron Bridge performs on-the-fly, model-parallel-aware, per-parameter conversion—unlike traditional converters that require a single GPU and full in-memory loading of both Megatron-Core and HF models. + +- For API-centric usage, see the guide: [Bridge with 🤗 Hugging Face](/./bridge-guide) + +## Architecture at a glance + +```mermaid +flowchart LR + HF[Hugging Face
PreTrained Model/Config] -->|detect| AutoBridge + AutoBridge -->|select| Bridge[MegatronModelBridge
architecture-specific] + Bridge -->|provider bridge| Provider[Model Provider
TransformerConfig builder] + Provider -->|instantiate| Megatron[Distributed Megatron Model] + + Bridge -->|mapping registry| Registry[MegatronMappingRegistry] + Registry -->|maps to| Mappings[Param Mapping
Auto Row Col QKV ...] + + HF <-->|convert per-parameter| Bridge + Megatron --|TP/PP/VPP-aware|--> Mappings +``` + +Key components: + +- AutoBridge: Detects HF architecture, constructs the appropriate bridge, exposes high-level convert/save APIs. See [API Documentation](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/). +- MegatronModelBridge: Orchestrates conversion, builds conversion tasks, handles per-parameter streaming. See [model_bridge.py](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/conversion/model_bridge.py). +- MegatronMappingRegistry: Registry of parameter-name mappings; resolves concrete `MegatronParamMapping` for each weight. See [mapping_registry.py](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/conversion/mapping_registry.py). +- Param Mappings: Implement parameter transformations and parallel distribution (Auto, ColumnParallel, RowParallel, QKV, GatedMLP, Replicated, custom). See [param_mapping.py](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/conversion/param_mapping.py). +- Model Providers: Build `TransformerConfig`-compatible providers for Megatron-Core and instantiate distributed models. See [models/](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models). +- Specific model bridge definitions: Architecture-specific bridges live under their model folders, for example [LlamaBridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/llama/llama_bridge.py) and [Qwen3Bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/qwen/qwen3_bridge.py). + +## Conversion workflow + +The conversion engine is powered by two pieces: the architecture-specific Megatron model bridge and the parameter mappings. + +(1) Config mapping + model creation: You specify a config mapping and parameter mapping strategy (name patterns + mapping types). The bridge loads the HF config, translates it into a Megatron provider, and the provider instantiates a (possibly distributed) Megatron model. With TP/PP/EP enabled, each rank holds only a shard of the model. + +```python +from megatron.bridge import AutoBridge + +# Build bridge and instantiate Megatron model(s) +bridge = AutoBridge.from_hf_pretrained("meta-llama/Llama-3.2-1B") +provider = bridge.to_megatron_provider() +provider.finalize() +megatron_model = provider.provide_distributed_model(wrap_with_ddp=False) +``` + +(2) Gather all params: After the model is created, the bridge enumerates all named parameters and buffers across PP ranks. It then sorts them to produce a deterministic global order, ensuring every rank uses the same mapping order for collective operations during conversion. + +(3) Resolve mappings: Using the global Megatron parameter names, the bridge queries the mapping registry to resolve a concrete mapping per parameter. For example, in Qwen3, `decoder.layers.0.self_attention.linear_qkv.weight` matches a `QKVMapping` pattern. Resolution always starts from the Megatron name; HF names are derived via wildcard substitution. Only the referenced HF tensors are fetched from safetensors—the full HF model is never fully loaded. + +```python +from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry +from megatron.bridge.models.conversion.param_mapping import AutoMapping, QKVMapping + +registry = MegatronMappingRegistry( + AutoMapping( + megatron_param="decoder.layers.*.mlp.linear_fc2.weight", + hf_param="model.layers.*.mlp.down_proj.weight", + ), + QKVMapping( + megatron_param="decoder.layers.*.self_attention.linear_qkv.weight", + q="model.layers.*.self_attn.q_proj.weight", + k="model.layers.*.self_attn.k_proj.weight", + v="model.layers.*.self_attn.v_proj.weight", + ), +) +# Example: "decoder.layers.0.self_attention.linear_qkv.weight" → QKVMapping +``` + +(4) Create conversion tasks: The bridge pairs every Megatron parameter with its resolved mapping and associated metadata (owning module, tensor handle, parallelism context). These per-parameter tasks become the unit of work for conversion. + +(5) Execute conversions: For HF→Megatron or Megatron→HF, the bridge iterates tasks and calls the mapping’s `hf_to_megatron` or `megatron_to_hf` routines. Conversion streams parameter-by-parameter to minimize memory. + +```python +# HF → Megatron streaming import (internally iterates conversion tasks) +bridge.load_hf_weights(megatron_model) +``` + +(6) Mapping semantics: Each mapping handles the necessary distribution logic—broadcasting across PP, scattering/gathering across TP/EP—and applies structural transforms (e.g., QKV fuse/split, gated-MLP concat/split, row/column parallel splits). + +Properties: + +- Per-parameter streaming: Only the weight currently being processed is kept in memory. +- Parallelism aware: Distribution respects TP (tensor), PP (pipeline), VPP (virtual pipeline), and expert parallel settings. +- Deterministic mapping: Names resolve via `MegatronMappingRegistry`, including wildcard support. + +### HF → Megatron (import) + +```mermaid +sequenceDiagram + participant U as User + participant AB as AutoBridge + participant B as MegatronModelBridge + participant R as MappingRegistry + participant M as Megatron Model(s) + + U->>AB: from_hf_pretrained(path_or_id) + AB->>B: select bridge for architecture + U->>AB: load_hf_weights(model) + AB->>B: load_hf_weights(model) + B->>M: enumerate params/buffers across PP ranks + B->>B: gather global param names (all ranks) + B->>B: sort names deterministically + B->>R: mapping_registry() + B->>R: build_conversion_tasks(megatron_params, hf_state) + loop for each global Megatron param name + R-->>B: resolve mapping by Megatron name + B->>B: fetch referenced HF tensor(s) (safetensors) + B->>B: apply transforms (QKV/GatedMLP/Row/Column) + B->>M: scatter/gather TP/EP, broadcast PP + B->>M: copy into destination tensor + end + B-->>U: model with loaded weights +``` + +### Megatron → HF (export) + +```mermaid +sequenceDiagram + participant U as User + participant AB as AutoBridge + participant B as MegatronModelBridge + participant R as MappingRegistry + participant M as Megatron Model(s) + + U->>AB: export_hf_weights(model) / save_hf_pretrained(model, path) + AB->>B: export_hf_weights(model) + B->>M: enumerate params/buffers across PP ranks + B->>B: gather global param names (all ranks) + B->>B: sort names deterministically + B->>R: mapping_registry() + B->>R: build_conversion_tasks(megatron_params) + loop for each global Megatron param name + R-->>B: resolve mapping by Megatron name + B->>M: read shard(s) from owning PP rank + B->>M: gather across TP/EP as needed + B->>B: apply inverse transforms (split QKV/GatedMLP/Row/Column) + alt streaming + B-->>U: yield (hf_name, tensor) + else saving + B->>U: write to safetensors / HF dir + end + end + B-->>U: export complete +``` + +## Param mappings and parallelism + +Mapping types available via [param_mapping.py](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/conversion/param_mapping.py): + +- AutoMapping: General-purpose 1:1 parameter mapping with automatic TP-type detection; dispatches to ColumnParallelMapping, RowParallelMapping, or ReplicatedMapping based on the layer/module type (wildcards supported). Participates in PP broadcast and EP gather when applicable. +- ColumnParallelMapping: Splits along the output dimension (dim 0) under TP. Participates in PP broadcast and EP gather when applicable. +- RowParallelMapping: Splits along the input dimension (dim 1) under TP. Participates in PP broadcast and EP gather when applicable. +- QKVMapping: Fuses/splits HF Q, K, V projections to Megatron's interleaved QKV format and vice versa. Uses PP broadcast as needed and delegates TP to the underlying mapping. +- GatedMLPMapping: Concatenates/splits gate and up projections. Participates in PP broadcast and EP gather when applicable. +- ReplicatedMapping: Keeps parameters fully replicated across TP ranks (e.g., LayerNorm). Participates in PP broadcast and EP gather when applicable. + +Note: If you need a one-to-many or many-to-one mapping that is not covered by QKVMapping or GatedMLPMapping, implement a custom mapping. + +### Example Mapping - ColumnParallelMapping: PP, TP, EP in practice + +- HF → Megatron (import): + - HF tensors are available from storage to all ranks; TP rank 0 reads the full tensor and performs the split/scatter. + - TP: Rank 0 splits along dim 0 into `tp_size` chunks and scatters shards to TP ranks so each rank receives a tensor matching its local parameter shape/dtype/device. + - PP: No PP collectives are needed; the owning PP stage writes its shard directly. + - EP: For expert parameters, each EP rank receives its local experts by name; no cross-EP collectives are required on import. + +- Megatron → HF (export): + - Only the owning PP stage initially holds the local Megatron shard; it broadcasts to all PP ranks before TP gather. + - PP: The owning PP stage first broadcasts the tensor to all PP ranks so every rank participates in the collectives. + - TP: All TP shards are gathered and concatenated along dim 0 to reconstruct the full tensor. + - EP: For expert parameters, shards are gathered across EP ranks and one HF tensor per expert is emitted with the correct names. + - Let total experts be E and EP size be S (assume E % S == 0). Each EP rank owns E/S experts. For a given local expert index L on each EP rank, the global expert ids are L, L+E/S, ..., L+(S-1)*E/S. We gather tensors from all EP ranks and emit one HF tensor per global expert id by substituting that id into the HF parameter name. + +This mirrors [ColumnParallelMapping.hf_to_megatron](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/conversion/param_mapping.py) and [ColumnParallelMapping.megatron_to_hf](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/conversion/param_mapping.py) in [param_mapping.py](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/conversion/param_mapping.py). + +Implementation notes (from code): +- Dtype handling: When HF and Megatron dtypes differ, weights are cast to the Megatron parameter dtype with a warning before TP scatter (see ColumnParallelMapping.hf_to_megatron in [param_mapping.py](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/conversion/param_mapping.py)). +- FP8 export: Tensors are dequantized on export when using FP8 tensor classes (see `maybe_dequantize` in [param_mapping.py](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/conversion/param_mapping.py)). +- MoE experts: Expert parameter names are normalized for lookup and expert shards are gathered across EP ranks and re-emitted per global expert id (see `gather_from_ep_ranks` in [param_mapping.py](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/conversion/param_mapping.py)). + +## Architecture-specific bridge example: Qwen3 + +Embedded from `src/megatron/bridge/models/qwen/qwen3_bridge.py`: + +```python +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch +from megatron.core.models.gpt.gpt_model import GPTModel +from transformers import Qwen3ForCausalLM + +from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry +from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge +from megatron.bridge.models.conversion.param_mapping import ( + AutoMapping, + GatedMLPMapping, + QKVMapping, +) + +@MegatronModelBridge.register_bridge(source=Qwen3ForCausalLM, target=GPTModel, model_type="qwen3") +class Qwen3Bridge(MegatronModelBridge): + """ + Megatron Bridge for Qwen3 Causal LM. + + This bridge handles the conversion between HuggingFace Qwen3ForCausalLM + and Megatron-Core GPTModel formats. Qwen3 differs from Qwen2 by using + QK layernorm and no QKV bias. + + Example: + >>> from megatron.bridge import AutoBridge + >>> bridge = AutoBridge.from_hf_pretrained("Qwen/Qwen3-1.7B") + >>> provider = bridge.to_megatron_provider() + """ + + def provider_bridge(self, hf_pretrained): + """Convert HuggingFace Qwen3 config to GPTModelProvider.""" + provider = super().provider_bridge(hf_pretrained) + + provider.normalization = "RMSNorm" + provider.gated_linear_unit = True + provider.position_embedding_type = "rope" + provider.add_bias_linear = False + provider.add_qkv_bias = False # Qwen3 does NOT have QKV bias (unlike Qwen2) + provider.hidden_dropout = 0.0 + provider.qk_layernorm = True # Qwen3 uses QK layernorm + provider.autocast_dtype = torch.bfloat16 + + return provider + + def mapping_registry(self) -> MegatronMappingRegistry: + # Return MegatronMappingRegistry containing parameter mappings from Megatron to HF format + # First create simple 1:1 parameter mappings using a dictionary for readability + + # Dictionary maps Megatron parameter names -> HF parameter names + # Supports wildcard (*) patterns for layer-specific parameters + param_mappings = { + "embedding.word_embeddings.weight": "model.embed_tokens.weight", + "output_layer.weight": "lm_head.weight", + "decoder.final_layernorm.weight": "model.norm.weight", + "decoder.layers.*.self_attention.linear_qkv.layer_norm_weight": "model.layers.*.input_layernorm.weight", + "decoder.layers.*.mlp.linear_fc1.layer_norm_weight": "model.layers.*.post_attention_layernorm.weight", + "decoder.layers.*.self_attention.q_layernorm.weight": "model.layers.*.self_attn.q_norm.weight", # Qwen3 specific + "decoder.layers.*.self_attention.k_layernorm.weight": "model.layers.*.self_attn.k_norm.weight", # Qwen3 specific + "decoder.layers.*.self_attention.linear_proj.weight": "model.layers.*.self_attn.o_proj.weight", + "decoder.layers.*.mlp.linear_fc2.weight": "model.layers.*.mlp.down_proj.weight", + } + + mapping_list = [] + # Convert each dictionary entry to AutoMapping(megatron_param, hf_param) + for megatron_param, hf_param in param_mappings.items(): + mapping_list.append(AutoMapping(megatron_param=megatron_param, hf_param=hf_param)) + + # Add special mappings that require parameter concatenation/transformation + mapping_list.extend( + [ + # QKV: Combine separate Q, K, V matrices into single QKV matrix + # Note: Qwen3 does NOT have bias in QKV projections (unlike Qwen2) + QKVMapping( + megatron_param="decoder.layers.*.self_attention.linear_qkv.weight", + q="model.layers.*.self_attn.q_proj.weight", + k="model.layers.*.self_attn.k_proj.weight", + v="model.layers.*.self_attn.v_proj.weight", + ), + # Gated MLP: Combine gate and up projection matrices into single FC1 matrix + GatedMLPMapping( + megatron_param="decoder.layers.*.mlp.linear_fc1.weight", + gate="model.layers.*.mlp.gate_proj.weight", + up="model.layers.*.mlp.up_proj.weight", + ), + ] + ) + + return MegatronMappingRegistry(*mapping_list) + +``` + +Notes: + +- `provider_bridge`: Translate HF config into a Megatron-compatible provider, including architecture quirks (e.g., `qk_layernorm=True`). +- `mapping_registry`: Define exact name patterns and transformation mappings. Wildcards `*` apply the same rule across layers. diff --git a/fern/v0.2.0/pages/documentation.mdx b/fern/v0.2.0/pages/documentation.mdx new file mode 100644 index 0000000000..238523af36 --- /dev/null +++ b/fern/v0.2.0/pages/documentation.mdx @@ -0,0 +1,139 @@ +--- +title: "Documentation Development" +description: "" +--- + +## Build the Documentation + +The following sections describe how to set up and build the Megatron Bridge documentation. + +Switch to the documentation source folder and generate HTML output. + +```sh +cd docs/ +uv run --only-group docs sphinx-build . _build/html +``` + +* The resulting HTML files are generated in a `_build/html` folder that is created under the project `docs/` folder. +* The generated python API docs are placed in `apidocs` under the `docs/` folder. + + +If you encounter the error "Failed to generate package metadata for megatron-core @ directory+3rdparty/Megatron-LM," +run the command to install the necessary submodules: + +`git submodule update --init --recursive` + + +## Live Building + +When writing documentation, it can be helpful to serve the documentation and have it update live while you edit. + +To do so, run: + +```sh +cd docs/ +uv run --only-group docs sphinx-autobuild . _build/html --port 12345 --host 0.0.0.0 +``` + +Open a web browser and go to `http://${HOST_WHERE_SPHINX_COMMAND_RUN}:12345` to view the output. + +## Write Tests in Python Docstrings + +Any code in triple backtick blocks with the `{doctest}` directive will be tested. The format follows Python's doctest module syntax, where `>>>` indicates Python input and the following line shows the expected output. Here's an example: + +```python +def add(x: int, y: int) -> int: + """ + Adds two integers together. + + Args: + x (int): The first integer to add. + y (int): The second integer to add. + + Returns: + int: The sum of x and y. + + Examples: + ```python +>>> from megatron.bridge.made_up_package import add + >>> add(1, 2) + 3 +``` + + """ + return x + y +``` + +## Run Tests in Python Docstrings + +You can run tests in our Python docstrings with: + +```sh +cd docs/ +uv run --only-group docs sphinx-build -b doctest . _build/doctest +``` + +## Documentation Version + +The three files below control the version switcher. Before you attempt to publish a new version of the documentation, update these files in the docs/ folder to match the latest version numbers. + +The ``version`` and ``release`` variables with your GitHub releases when publishing new versions of documentation. + + +Latest should only be ``version`` and ``release`` variables in the main branch. + + +### versions1.json + +This JSON file defines the versions displayed in the switcher drop down. When adding a new version to the JSON, please make sure the ``version`` and ``url`` contain same version as your release. + +Example: + +```json +[ + { + "preferred": true, + "version": "latest", + "url": "https://docs.nvidia.com/nemo/megatron-bridge/latest/" + }, + { + + "version": "#.#.#", + "url": "https://docs.nvidia.com/nemo/megatron-bridge/#.#.#/" + }, + { + + "version": "0.1.0", + "url": "https://docs.nvidia.com/nemo/megatron-bridge/0.1.0/" + } +] +``` + + +Use absolute URLs for the ``url`` variable. + + +## project.json + +This JSON file tells the version switcher that documentation matches the selected version in the switcher. The ``version`` should contain same version as your release. + +```json +{ + "name": "megatron-bridge", + "version": "#.#.#" +} +``` + +## conf.py + +The conf.py ``release`` should contain same version as your release. + +```python +# -- Project information ----------------------------------------------------- +# https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information + +project = "Megatron Bridge" +copyright = "2025, NVIDIA Corporation" +author = "NVIDIA Corporation" +release = "#.#.#" +``` diff --git a/fern/v0.2.0/pages/index.mdx b/fern/v0.2.0/pages/index.mdx new file mode 100644 index 0000000000..5f5030fb6b --- /dev/null +++ b/fern/v0.2.0/pages/index.mdx @@ -0,0 +1,223 @@ +--- +title: "NeMo Megatron Bridge" +description: "" +--- + +
+ +# NeMo Megatron Bridge + +[![codecov](https://codecov.io/github/NVIDIA-NeMo/Megatron-Bridge/graph/badge.svg?token=4NMKZVOW2Z)](https://codecov.io/github/NVIDIA-NeMo/Megatron-Bridge) +[![CICD NeMo](https://github.com/NVIDIA-NeMo/Megatron-Bridge/actions/workflows/cicd-main.yml/badge.svg)](https://github.com/NVIDIA-NeMo/Megatron-Bridge/actions/workflows/cicd-main.yml) +[![Python 3.10+](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/release/python-3100/) +[![GitHub Stars](https://img.shields.io/github/stars/NVIDIA-NeMo/Megatron-Bridge.svg?style=social&label=Star&cacheSeconds=14400)](https://github.com/NVIDIA-NeMo/Megatron-Bridge/stargazers/) + +[Documentation](https://docs.nvidia.com/nemo/megatron-bridge/latest/) | [Supported Models](#supported-models) | [Examples](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/examples) | [Contributing](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/CONTRIBUTING.md) +
+ +## 📣 News +* [12/16/2025] [Mind Lab](https://macaron.im/mindlab) successfully used Megatron-bridge and [VeRL](https://github.com/volcengine/verl) to trained GRPO Lora for Trillion-parameter model on 64 H800 - See their [techblog](https://macaron.im/mindlab/research/building-trillion-parameter-reasoning-rl-with-10-gpus). +* [12/15/2025] Day 0 support for [NVIDIA-NeMotron-3-Nano-30B-A3B-FP8](https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8)! [Reproducible code](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/nano-v3) and custom NGC container: [nvcr.io/nvidia/nemo:25.11.nemotron_3_nano](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/nemo?version=25.11.nemotron_3_nano) + +## Overview + +NeMo Megatron Bridge is a PyTorch-native library within the [NeMo Framework](https://github.com/NVIDIA-NeMo) that provides pretraining, SFT and LoRA for popular LLM and VLM models. It serves as a powerful **bridge, conversion, and verification layer** between 🤗 Hugging Face and [Megatron Core](https://github.com/NVIDIA/Megatron-LM/tree/main/megatron/core). It provides bidirectional checkpoint conversion between these formats, enabling other projects to leverage Megatron Core's parallelism capabilities or export models for various inference engines. The bridge includes built-in verification mechanisms to ensure conversion accuracy and checkpoint integrity across different model formats. + +On top of the bridge, NeMo Megatron Bridge provides a performant and scalable PyTorch-native training loop that leverages [Megatron Core](https://github.com/NVIDIA/Megatron-LM/tree/main/megatron/core) to deliver state-of-the-art training throughput. It supports pretraining and fine-tuning with features like tensor and pipeline parallelism, and mixed precision (FP8, BF16, FP4, etc.). Users can either use existing 🤗 Hugging Face models or define custom PyTorch model definitions for flexible end-to-end workflows. + +NeMo Megatron Bridge is a refactor of the [previous NeMo](https://github.com/NVIDIA/NeMo) training stack that adopts a PyTorch-native training loop to provide greater flexibility and customizability for developers. + +![image](/assets/Repo-Mbridge.png) + +## 🔧 Installation + +### 🐳 NeMo Framework container + +The best experience, highest performance, and full feature support are provided by the [NeMo Framework container](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/nemo/tags). Fetch the most recent $TAG and run the following to start a container: + +```bash +docker run --rm -it -w /workdir -v $(pwd):/workdir \ + --entrypoint bash \ + --gpus all \ + nvcr.io/nvidia/nemo:${TAG} +``` + +For development installation and additional details, please refer to our [Contribution guide](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/CONTRIBUTING.md). + +## ⚡ Quickstart + +To get started, install Megatron Bridge or download a NeMo Framework container as described [above](#-installation). + +Log in to Hugging Face Hub: + +```sh +huggingface-cli login --token +``` + +Conversion-only quickstart (✅ Core): + +```python +from megatron.bridge import AutoBridge + +# 1) Create a bridge from a Hugging Face model (hub or local path) +bridge = AutoBridge.from_hf_pretrained("meta-llama/Llama-3.2-1B", trust_remote_code=True) + +# 2) Get a Megatron provider and configure parallelism before instantiation +provider = bridge.to_megatron_provider() +provider.tensor_model_parallel_size = 1 +provider.pipeline_model_parallel_size = 1 +provider.finalize() +# 3) Materialize Megatron Core model(s) +model = provider.provide_distributed_model(wrap_with_ddp=False) + +# 4a) Export Megatron → Hugging Face (full HF folder with config/tokenizer/weights) +bridge.save_hf_pretrained(model, "./hf_exports/llama32_1b") + +# 4b) Or stream only weights (Megatron → HF) +for name, weight in bridge.export_hf_weights(model, cpu=True): + print(name, tuple(weight.shape)) +``` + +Training quickstart using pre-configured recipes: + +```python +from megatron.bridge.recipes.llama import llama32_1b_pretrain_config +from megatron.bridge.training.gpt_step import forward_step +from megatron.bridge.training.pretrain import pretrain + +if __name__ == "__main__": + # The recipe uses the Llama 3.2 1B model configuration from HuggingFace + cfg = llama32_1b_pretrain_config(seq_length=1024) + + # Override training parameters + cfg.train.train_iters = 10 + cfg.scheduler.lr_decay_iters = 10000 + cfg.model.vocab_size = 8192 + cfg.tokenizer.vocab_size = cfg.model.vocab_size + + pretrain(cfg, forward_step) +``` + +You can launch the above script with: + +```sh +torchrun --nproc-per-node= /path/to/script.py +``` + +More examples: + +- [Conversion scripts overview](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/conversion/README.md) +- [Import/Export checkpoints](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/conversion/convert_checkpoints.py) +- [Generation with bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/conversion/hf_to_megatron_generate_text.py) +- [Multi-GPU loading from HF](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/conversion/hf_megatron_roundtrip_multi_gpu.py) +- [Compare HF vs Megatron outputs](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/conversion/compare_models.py) +- [Toy RLHF with Bridge (HF inference + Megatron training)](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/rl/rlhf_with_bridge.py) + +For a deeper dive into conversion design and advanced usage, see the [models README](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/models/README.md). + +## 🚀 Key Features + +- **Bridge with 🤗 Hugging Face**: Seamless bidirectional conversion between 🤗 Hugging Face and Megatron formats for interoperability ([model bridges](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models), [auto bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/models/conversion/auto_bridge.py), [conversion examples](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/examples/conversion)) + - Online import/export without intermediate full checkpoints + - Parallelism-aware (TP/PP/VPP/CP/EP/ETP) during conversion + - Memory-efficient per-parameter streaming + - Simple high-level `AutoBridge` API with architecture auto-detection + - Optimized paths when Transformer Engine is available +- **Flexible to Customize**: Lightweight custom training loop making it easy to configure custom logic in data loading, distributed training, checkpointing, evaluation and logging ([training framework](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/training), [training utilities](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/training/utils)) +- **Supervised & Parameter-Efficient Finetuning**: SFT & PEFT implementation tailored for Megatron-based models that supports LoRA, DoRA, and user-defined PEFT methods ([PEFT implementations](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/peft), [finetune module](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/training/finetune.py), [SFT dataset](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/data/datasets/sft.py)) +- **SOTA Training Recipes**: Pre-configured production-ready training recipes for popular models like Llama 3, with optimized hyperparameters and distributed training configuration ([Llama recipes](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/recipes/llama), [recipe examples](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/examples/models)) +- **Performance Optimization**: Built-in support for FP8 training, model parallelism, and memory-efficient techniques to offer high utilization and near-linear scalability to thousands of nodes. ([mixed precision](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/training/mixed_precision.py), [communication overlap](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/training/comm_overlap.py), [optimizer utilities](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/recipes/utils/optimizer_utils.py)) + +## Supported Models + +Megatron Bridge provides out-of-the-box bridges and training recipes for a wide range of models, built on top of base model architectures from [Megatron Core](https://github.com/NVIDIA/Megatron-LM/tree/main/megatron/core). Refer to the [models directory](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models) for the most up-to-date list of model bridges. + +### Supported Models Overview + +For more details on supported models, see our documentation: +- **[Large Language Models](https://docs.nvidia.com/nemo/megatron-bridge/latest/models/llm/index.html)** +- **[Vision Language Models](https://docs.nvidia.com/nemo/megatron-bridge/latest/models/vlm/index.html)** + +| Model | Checkpoint Conversion | Pretrain Recipes | SFT & LoRA Recipes | +|-------|-------------------|-------------------|-------------------| +| [DeepSeek V2](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/deepseek) | ✅ | ✅ ([v2](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/recipes/deepseek/deepseek_v2.py)) | Coming soon | +| [DeepSeek V2 Lite](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/deepseek) | ✅ | ✅ ([v2-lite](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/recipes/deepseek/deepseek_v2.py)) | Coming soon | +| [DeepSeek V3](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/deepseek) | ✅ | ✅ ([v3](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/recipes/deepseek/deepseek_v3.py)) | Coming soon | +| [Gemma](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/gemma) | ✅ | Coming soon | Coming soon | +| [Gemma 2](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/gemma) | ✅ | Coming soon | Coming soon | +| [Gemma 3](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/gemma) | ✅ | ✅ ([1B](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/recipes/gemma/gemma3.py)) | ✅ ([1B](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/recipes/gemma/gemma3.py)) | +| [Gemma 3-VL](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/gemma_vl) | ✅ | Coming soon | ✅ ([4B/12B/27B](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/recipes/gemma3_vl/gemma3_vl.py)) | +| [GLM-4.5](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/glm) | ✅ | ✅ ([106B-Air/355B](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/recipes/glm/glm45.py)) | ✅ ([106B-Air/355B](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/recipes/glm/glm45.py)) | +| [GPT-oss](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/gpt_oss) | ✅ | ✅ ([20B/120B](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/recipes/gpt_oss/gpt_oss.py)) | ✅ ([20B/120B](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/recipes/gpt_oss/gpt_oss.py)) | +| [Llama 2](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/llama) | ✅ | ✅ ([7B](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/recipes/llama/llama2.py)) | Coming soon | +| [Llama 3](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/llama) | ✅ | ✅ ([8B/70B](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/recipes/llama/llama3.py)) | ✅ ([8B/70B](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/recipes/llama/llama3.py)) | +| [Llama 3.1](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/llama) | ✅ | ✅ ([8B/70B/405B](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/recipes/llama/llama3.py)) | ✅ ([8B/70B/405B](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/recipes/llama/llama3.py)) | +| [Llama 3.2](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/llama) | ✅ | ✅ ([1B/3B](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/recipes/llama/llama3.py)) | ✅ ([1B/3B](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/recipes/llama/llama3.py)) | +| [Llama 3.3](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/llama) | ✅ | Coming soon | Coming soon | +| [Llama Nemotron](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/llama_nemotron) | ✅ | Coming soon | Coming soon | +| [Mistral](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/mistral) | ✅ | Coming soon | Coming soon | +|[Ministral](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/ministral3)| ✅| ✅ [3B/8B/14B](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/recipes/ministral3/ministral3.py)|✅ [3B/8B/14B](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/recipes/ministral3/ministral3.py)| +| [Moonlight](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/deepseek) | ✅ | ✅ ([16B](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/recipes/moonlight/moonlight_16b.py)) | ✅ ([16B](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/recipes/moonlight/moonlight_16b.py)) | +| [Nemotron](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/nemotron) | ✅ | Coming soon | Coming soon | +| [Nemotron-3](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/nano-v3/src/megatron/bridge/models/nemotronh) | ✅ | ✅ ([A3B](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/nano-v3/src/megatron/bridge/recipes/nemotronh/nemotron_3_nano.py)) | ✅ ([A3B](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/nano-v3/src/megatron/bridge/recipes/nemotronh/nemotron_3_nano.py)) | +| [Nemotron-H](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/nemotronh) | ✅ | ✅ ([4B/8B/47B/56B](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/recipes/nemotronh/nemotronh.py)) | Coming soon | +| [Nemotron Nano v2](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/nemotronh) | ✅ | ✅ ([9B/12B](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/recipes/nemotronh/nemotron_nano_v2.py)) | Coming soon | +| [Nemotron Nano v2 VL](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/nemotron_vl) | ✅ | Coming soon | ✅ ([9B/12B](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/recipes/nemotron_vl/nemotron_nano_v2_vl.py)) | +| [OlMoE](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/olmoe) | ✅ | ✅ ([7B](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/recipes/olmoe/olmoe_7b.py)) | ✅ ([7B](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/recipes/olmoe/olmoe_7b.py)) | +| [Qwen2](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/qwen) | ✅ | ✅ ([500M/1.5B/7B/72B](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/recipes/qwen/qwen2.py)) | ✅ ([500M/1.5B/7B/72B](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/recipes/qwen/qwen2.py)) | +| [Qwen2.5](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/qwen) | ✅ | ✅ ([500M/1.5B/7B/14B/32B/72B](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/recipes/qwen/qwen2.py)) | ✅ ([500M/1.5B/7B/14B/32B/72B](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/recipes/qwen/qwen2.py)) | +| [Qwen2.5-VL](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/qwen_vl) | ✅ | Coming soon | ✅ ([3B/7B/32B/72B](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/recipes/qwen_vl/qwen25_vl.py)) | +| [Qwen3](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/qwen) | ✅ | ✅ ([600M/1.7B/4B/8B/14B/32B](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/recipes/qwen/qwen3.py)) | ✅ ([600M/1.7B/4B/8B/14B/32B](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/recipes/qwen/qwen3.py)) | +| [Qwen3-MoE](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/qwen) | ✅ | ✅ ([A3B/A22B](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/recipes/qwen/qwen3_moe.py)) | ✅ ([A3B/A22B](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/recipes/qwen/qwen3_moe.py)) | +| [Qwen3 Next](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/qwen) | ✅ | ✅ ([80B-A3B](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/recipes/qwen/qwen3_next.py)) | ✅ ([80B-A3B](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/recipes/qwen/qwen3_next.py)) | +| [Qwen3-VL](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/qwen_vl) | ✅ | Coming soon | ✅ ([8B/A3B-A30B-MoE](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/recipes/qwen_vl/qwen3vl.py)) | + +#### Launching Recipes + +For a conceptual overview of how recipes are structured, overridden, and launched with either `torchrun` or NeMo-Run, read the [Using Recipes guide](https://docs.nvidia.com/nemo/megatron-bridge/latest/recipe-usage.html). + +Runnable tutorials live in `tutorials/recipes/llama` that covers: + +- `00_quickstart_pretrain.py` for mock-data pretraining +- `01_quickstart_finetune.py` + LoRA configs +- YAML-driven flows and launch helpers + +{/* ### Vision-Language Models */} + +## Performance Benchmarks + +For detailed performance benchmarks including throughput metrics across different GPU systems (DGX-GB200, DGX-B200, DGX-H100) and model configurations, see the [Performance Summary](https://docs.nvidia.com/nemo/megatron-bridge/latest/performance-summary.html) in our documentation. + +## Project Structure + +``` +Megatron-Bridge/ +├── examples/ +│ ├── models/ # Bridge usage examples +│ └── recipes/ # Training examples +├── src/megatron/bridge/ +│ ├── data/ # Dataloaders and iterators +│ ├── models/ # Hugging Face bridge infrastructure and model-specific implementations +│ │ ├── llama/ # Llama model providers +│ │ └── .../ # Other models (gpt, t5, etc.) +│ ├── peft/ # PEFT transformations and wrappers +│ ├── recipes/ # Complete training recipes +│ ├── training/ # Training loop components +│ │ ├── tokenizers/ # Tokenizer library +│ │ └── utils/ # Training-specific utilities +│ └── utils/ # Generic utilities for repo-wide usage +└── tests/ # Comprehensive test suite +``` + +## Acknowledgement & Contributing + +Megatron-Bridge is the continuation of [MBridge](https://github.com/ISEEKYAN/mbridge) by [Yan Bai](https://github.com/ISEEKYAN). We appreciate all the contribution and adoptions by the community partners: + +- [Mind Lab](https://macaron.im/mindlab) successfully used Megatron-bridge and [VeRL](https://github.com/volcengine/verl) to trained GRPO Lora for Trillion-parameter model on 64 H800 - See their [techblog](https://macaron.im/mindlab/research/building-trillion-parameter-reasoning-rl-with-10-gpus). +- [VeRL](https://github.com/volcengine/verl) has adopted Megatron-Bridge as a connector to Megatron-Core and for LoRA support. +- [Slime](https://github.com/THUDM/slime) has adopted Megatron-Bridge as Megatron-Core checkpoint converter. +- [SkyRL](https://github.com/NovaSky-AI/SkyRL) has adopted Megatron-Bridge as Megatron-Core connector. +- [Nemo-RL](https://github.com/NVIDIA/nemo-rl) has adopted Megatron-Bridge as Megatron-Core connector. +- Community contributions: Special thanks to [Guanyou He](https://github.com/Thaurun) and [Junyu Wu](https://github.com/nrailg) from Weixin Group Infrastructure Center. + +Please see our [Contributor Guidelines](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/CONTRIBUTING.md) for more information on how to get involved. diff --git a/fern/v0.2.0/pages/megatron-lm-to-megatron-bridge.mdx b/fern/v0.2.0/pages/megatron-lm-to-megatron-bridge.mdx new file mode 100644 index 0000000000..2c8d6b86f8 --- /dev/null +++ b/fern/v0.2.0/pages/megatron-lm-to-megatron-bridge.mdx @@ -0,0 +1,295 @@ +--- +title: "Megatron-LM to Megatron Bridge Guide" +description: "" +--- + +Megatron Bridge is Python-first: configure models, data, and training via typed Python APIs. All configuration lives in a structured `ConfigContainer` (see [Configuration overview](/training/config-container-overview)). Any field can be overridden from the command line using Hydra/OmegaConf syntax in the example training scripts. + +## Quick start + +Run your example training entrypoint and override config keys directly: + +```bash +python examples/models/llama/pretrain_llama3_8b.py \ + train.micro_batch_size=2 \ + train.global_batch_size=128 \ + model.num_layers=32 model.hidden_size=4096 model.num_attention_heads=32 \ + model.max_position_embeddings=4096 \ + dataset.sequence_length=4096 \ + checkpoint.save=/workspace/ckpts checkpoint.save_interval=1000 \ + logger.wandb_project=my_proj logger.wandb_exp_name=exp1 +``` + +Notes: +- Config groups are nested: `rng`, `train`, `model`, `optimizer`, `ddp`, `scheduler`, `dataset`, `logger`, `tokenizer`, `checkpoint`, `dist`, `profiling`, `peft`, `comm_overlap`, `mixed_precision`, `inprocess_restart`. +- After overrides are applied, runtime validation computes any dependent fields (e.g., data-parallel size, scheduler steps) and checks consistency. + +## Mapping Megatron-LM arguments to Megatron Bridge config + +Below is a concise mapping from common `megatron-lm/megatron/training/arguments.py` flags to the new dataclass fields. If a field is not listed here (e.g., highly model-specific knobs), it typically lives under `model.*`, `optimizer.*`, `dataset.*`, or `tokenizer.*` with similar names. + +### Model topology and parallelisms + +| megatron-lm arguments | Megatron Bridge config | Description | +| --- | --- | --- | +| `--tensor-model-parallel-size` | `model.tensor_model_parallel_size` | TP degree. | +| `--pipeline-model-parallel-size` | `model.pipeline_model_parallel_size` | PP degree. | +| `--context-parallel-size` | `model.context_parallel_size` | CP degree. | +| `--expert-model-parallel-size` | `model.expert_model_parallel_size` | EP degree. | +| `--expert-tensor-parallel-size` | `model.expert_tensor_parallel_size` | Expert TP degree. | +| `--sequence-parallel` | `model.sequence_parallel` | Enable sequence parallelism. | +| `--account-for-embedding-in-pipeline-split` | `model.account_for_embedding_in_pipeline_split` | Asymmetric PP: embeddings. | +| `--account-for-loss-in-pipeline-split` | `model.account_for_loss_in_pipeline_split` | Asymmetric PP: loss. | + +### Model architecture knobs + +| megatron-lm arguments | Megatron Bridge config | Description | +| --- | --- | --- | +| `--untie-embeddings-and-output-weights` | `model.share_embeddings_and_output_weights=false` | Untie embeddings/output. | +| `--position-embedding-type` | `model.position_embedding_type` | `learned_absolute` or `rope`. | +| `--rotary-percent` | `model.rotary_percent` | Fraction of rotary dims. | +| `--rotary-base` | `model.rotary_base` | RoPE base. | +| `--rotary-seq-len-interpolation-factor` | `model.seq_len_interpolation_factor` | RoPE interpolation factor. | +| `--normalization` | `model.normalization` | LayerNorm/RMSNorm, etc. | +| `--swiglu` | `model.gated_linear_unit=true` | Enable SwiGLU MLP. | +| `--norm-epsilon` | `model.layernorm_epsilon` | Epsilon for norm layers. | +| `--num-layers` | `model.num_layers` | Number of transformer layers. | +| `--hidden-size` | `model.hidden_size` | Model hidden size. | +| `--ffn-hidden-size` | `model.ffn_hidden_size` | MLP expansion size. | +| `--num-attention-heads` | `model.num_attention_heads` | Attention heads. | +| `--kv-channels` | `model.kv_channels` | Key/Value channels per head. | +| `--group-query-attention` | `model.num_query_groups` | Set groups (enable GQA). | +| `--num-query-groups` | `model.num_query_groups` | Number of query groups. | +| `--qk-layernorm` | `model.qk_layernorm` | Enable QK LayerNorm. | +| `--seq-length` | `model.seq_length` | Max model sequence length. | +| `--max-position-embeddings` | `model.seq_length` | Alias used by HF conversions. | +| `--make-vocab-size-divisible-by` | `model.make_vocab_size_divisible_by` | TP padding multiple. | +| `--disable-bias-linear` | `model.add_bias_linear=false` | Disable linear bias. | +| `--use-flash-attn` | `model.attention_backend=flash` | Use FlashAttention backend. | +| `--init-method-std` | `model.init_method_std` | Weight init standard deviation. | +| `--attention-dropout` | `model.attention_dropout` | Attention dropout. | +| `--hidden-dropout` | `model.hidden_dropout` | Hidden dropout. | + +### MoE + +| megatron-lm arguments | Megatron Bridge config | Description | +| --- | --- | --- | +| `--num-experts` | `model.num_moe_experts` | Experts per MoE layer. | +| `--moe-ffn-hidden-size` | `model.moe_ffn_hidden_size` | Expert MLP hidden size. | +| `--moe-router-load-balancing-type` | `model.moe_router_load_balancing_type` | e.g., aux_loss or seq_aux_loss. | +| `--moe-router-topk` | `model.moe_router_topk` | Top-k experts per token. | +| `--moe-router-pre-softmax` | `model.moe_router_pre_softmax` | Pre-softmax routing. | +| `--moe-grouped-gemm` | `model.moe_grouped_gemm` | Grouped GEMM for MoE. | +| `--moe-aux-loss-coeff` | `model.moe_aux_loss_coeff` | Aux loss coefficient. | +| `--moe-token-dispatcher-type` | `model.moe_token_dispatcher_type` | Token dispatcher: alltoall or flex. | +| `--moe-flex-dispatcher-backend` | `model.moe_flex_dispatcher_backend` | MoE token dispatcher: deepep or hybridep | +| `--moe-permute-fusion` | `model.moe_permute_fusion` | Enable MoE permute fusion. | +| `--moe-router-fusion` | `model.moe_router_fusion` | Enable MoE router fusion. | +| `--moe-router-dtype` | `model.moe_router_dtype` | Router dtype (e.g., fp32). | + +### Mixed precision + +| megatron-lm arguments | Megatron Bridge config | Description | +| --- | --- | --- | +| `--bf16` | `mixed_precision` preset (e.g., "bf16_mixed") | Select a mixed-precision recipe; sets `model.bf16`/`optimizer.bf16`. | + +Mixed precision is selected via the `mixed_precision` config key (e.g., preset names like `bf16_mixed`, `bf16`, or `fp16`, depending on your codebase) and is applied to `model`, `optimizer`, and `ddp` during `runtime_config_update`. + +### Training + +| megatron-lm arguments | Megatron Bridge config | Description | +| --- | --- | --- | +| `--micro-batch-size` | `train.micro_batch_size` | Per-rank batch size before gradient accumulation. | +| `--global-batch-size` | `train.global_batch_size` | Total batch across DP and micro-batches. | +| `--train-samples` | `train.train_samples` | Total training samples (sample-based mode). | +| `--rampup-batch-size` | `train.rampup_batch_size` | Start size, increment, and sample count for linear batch ramp-up. | +| `--decrease-batch-size-if-needed` | `train.decrease_batch_size_if_needed` | Adjust GBS to remain divisible when DP changes. | +| `--empty-unused-memory-level` | `train.empty_unused_memory_level` | PyTorch CUDA empty_cache cadence (0, 1, or 2). | +| `--check-weight-hash-across-dp-replicas-interval` | `train.check_weight_hash_across_dp_replicas_interval` | Interval to validate DP weight consistency. | +| `--train-iters` | `train.train_iters` | Number of training iterations. | +| `--exit-interval` | `train.exit_interval` | Exit when iteration % interval == 0. | +| `--exit-duration-in-mins` | `train.exit_duration_in_mins` | Exit after N minutes. | +| `--exit-signal-handler` | `train.exit_signal_handler` | Save and shut down on SIGTERM. | +| `--manual-gc` | `train.manual_gc` | Enable manual Python GC scheduling. | +| `--manual-gc-interval` | `train.manual_gc_interval` | Steps between manual GC runs. | +| `--no-manual-gc-eval` | `train.manual_gc_eval=false` | Disable GC at eval boundaries. | +| `--eval-iters` | `train.eval_iters` | Eval iterations per validation run. | +| `--eval-interval` | `train.eval_interval` | Steps between validations. | +| `--skip-train` | `train.skip_train` | Skip training loop (eval-only). | + +### Scheduler / Regularization + +| megatron-lm arguments | Megatron Bridge config | Description | +| --- | --- | --- | +| `--lr-decay-style` | `scheduler.lr_decay_style` | LR schedule: constant/linear/cosine/ISR/WSD. | +| `--lr-decay-iters` | `scheduler.lr_decay_iters` | Iterations over which to decay LR. | +| `--lr-wsd-decay-style` | `scheduler.lr_wsd_decay_style` | WSD anneal style. | +| `--lr-wsd-decay-iters` | `scheduler.lr_wsd_decay_iters` | Iterations for WSD anneal phase. | +| `--lr-warmup-fraction` | `scheduler.lr_warmup_fraction` | Warmup as fraction of decay span. | +| `--lr-warmup-iters` | `scheduler.lr_warmup_iters` | Warmup iterations (absolute). | +| `--lr-warmup-init` | `scheduler.lr_warmup_init` | Initial LR at start of warmup. | +| `--lr-decay-samples` | `scheduler.lr_decay_samples` | Samples over which to decay LR (sample-based training). | +| `--lr-warmup-samples` | `scheduler.lr_warmup_samples` | Warmup samples (sample-based training). | +| `--lr` | `optimizer.lr` | Base learning rate. | +| `--min-lr` | `optimizer.min_lr` | Minimum learning rate. | +| `--clip-grad` | `optimizer.clip_grad` | Gradient clipping value. | +| `--weight-decay` | `optimizer.weight_decay` | Weight decay. | +| `--adam-beta1` | `optimizer.adam_beta1` | Adam beta1. | +| `--adam-beta2` | `optimizer.adam_beta2` | Adam beta2. | +| `--override-opt_param-scheduler` | `scheduler.override_opt_param_scheduler` | Ignore ckpt scheduler and use config. | +| `--use-checkpoint-opt_param-scheduler` | `scheduler.use_checkpoint_opt_param_scheduler` | Load scheduler from checkpoint. | +| `--start-weight-decay` | `scheduler.start_weight_decay` | WD at start (non-constant modes). | +| `--end-weight-decay` | `scheduler.end_weight_decay` | WD at end (non-constant modes). | +| `--weight-decay-incr-style` | `scheduler.weight_decay_incr_style` | WD schedule: constant/linear/cosine. | + +### Checkpointing + +| megatron-lm arguments | Megatron Bridge config | Description | +| --- | --- | --- | +| `--save` | `checkpoint.save` | Directory to write checkpoints. | +| `--save-interval` | `checkpoint.save_interval` | Iterations between persistent saves. | +| `--no-save-optim` | `checkpoint.save_optim=false` | Do not save optimizer state. | +| `--no-save-rng` | `checkpoint.save_rng=false` | Do not save RNG state. | +| `--load` | `checkpoint.load` | Directory to load from. | +| `--no-load-optim` | `checkpoint.load_optim=false` | Do not load optimizer state. | +| `--load-main-params-from-ckpt` | `checkpoint.load_main_params_from_ckpt` | Load FP32 main params directly. | +| `--no-load-rng` | `checkpoint.load_rng=false` | Do not load RNG state. | +| `--non-persistent-save-interval` | `checkpoint.non_persistent_save_interval` | Frequency for ephemeral saves. | +| `--non-persistent-ckpt-type` | `checkpoint.non_persistent_ckpt_type` | Kind of ephemeral checkpoint (global/local/memory). | +| `--non-persistent-global-ckpt-dir` | `checkpoint.non_persistent_global_ckpt_dir` | Dir for global ephemeral saves. | +| `--non-persistent-local-ckpt-dir` | `checkpoint.non_persistent_local_ckpt_dir` | Dir for local-per-rank ephemeral saves. | +| `--non-persistent-local-ckpt-algo` | `checkpoint.non_persistent_local_ckpt_algo` | Local save algorithm selection. | +| `--finetune` | `checkpoint.finetune` | Load weights, reset iters, no optim/rng. | +| `--pretrained-checkpoint` | `checkpoint.pretrained_checkpoint` | Path to pretrained weights for finetune/SFT. | +| `--ckpt-step` | `checkpoint.ckpt_step` | Explicit step to load. | +| `--use-checkpoint-args` | `checkpoint.use_checkpoint_args` | Override model args from checkpoint metadata. | +| `--exit-on-missing-checkpoint` | `checkpoint.exit_on_missing_checkpoint` | Exit if `load` not found. | +| `--ckpt-format` | `checkpoint.ckpt_format` | Format: torch_dist/zarr/fsdp_dtensor. | +| `--ckpt-convert-format` | `checkpoint.ckpt_convert_format` | Conversion target format. | +| `--ckpt-convert-save` | `checkpoint.ckpt_convert_save` | Output dir for converted ckpt. | +| `--no-ckpt-fully-parallel-save` | `checkpoint.fully_parallel_save=false` | Disable DP-parallel save. | +| `--async-save` | `checkpoint.async_save` | Enable async saves (torch_dist only). | +| `--use-persistent-ckpt-worker` | `checkpoint.use_persistent_ckpt_worker` | Background worker for async saves. | +| `--ckpt-fully-parallel-load` | `checkpoint.fully_parallel_load` | Enable DP-parallel load. | +| `--ckpt-assume-constant-structure` | `checkpoint.ckpt_assume_constant_structure` | Optimize for fixed structure. | +| `--dist-ckpt-strictness` | `checkpoint.dist_ckpt_strictness` | Handling of key mismatches on load. | +| `--auto-detect-ckpt-format` | `checkpoint.auto_detect_ckpt_format` | Auto-detect checkpoint format on load. | +| `--replication` | `checkpoint.replication` | Enable replication of local checkpoints. | +| `--replication-jump` | `checkpoint.replication_jump` | Spacing between replica ranks. | +| `--replication-factor` | `checkpoint.replication_factor` | Number of replicas. | +| `--no-strict-fsdp-dtensor-load` | `checkpoint.strict_fsdp_dtensor_load=false` | Relax FSDP-DTensor strict load. | + +### Logging + +| megatron-lm arguments | Megatron Bridge config | Description | +| --- | --- | --- | +| `--log-interval` | `logger.log_interval` | Steps between console logs. | +| `--log-params-norm` | `logger.log_params_norm` | Compute and log parameter L2 norm. | +| `--log-throughput` | `logger.log_throughput` | Log tokens/sec per GPU. | +| `--log-progress` | `logger.log_progress` | Write progress.txt with tokens and FLOPs. | +| `--timing-log-level` | `logger.timing_log_level` | 0=min; 1=coarse ops; 2=many ops. | +| `--timing-log-option` | `logger.timing_log_option` | max/minmax/all across ranks. | +| `--tensorboard-dir` | `logger.tensorboard_dir` | TensorBoard log directory. | +| `--tensorboard-log-interval` | `logger.tensorboard_log_interval` | Steps between TB events. | +| `--tensorboard-queue-size` | `logger.tensorboard_queue_size` | Pending TB event queue size. | +| `--log-timers-to-tensorboard` | `logger.log_timers_to_tensorboard` | Write timers to TB. | +| `--no-log-loss-scale-to-tensorboard` | `logger.log_loss_scale_to_tensorboard=false` | Disable loss-scale TB logs. | +| `--log-validation-ppl-to-tensorboard` | `logger.log_validation_ppl_to_tensorboard` | Write validation perplexity (ppl) to TB. | +| `--log-memory-to-tensorboard` | `logger.log_memory_to_tensorboard` | Enable memory stats in TB. | +| `--log-world-size-to-tensorboard` | `logger.log_world_size_to_tensorboard` | Log world size in TB. | +| `--wandb-project` | `logger.wandb_project` | Weights & Biases project. | +| `--wandb-entity` | `logger.wandb_entity` | Weights & Biases entity/team. | +| `--wandb-exp-name` | `logger.wandb_exp_name` | Run name in W&B. | +| `--wandb-save-dir` | `logger.wandb_save_dir` | Local directory for W&B artifacts. | +| `--logging-level` | `logger.logging_level` | Python logging level (e.g., 20=INFO). | +| `--log-energy` | `logger.log_energy` | Log energy in Joules (if available). | + +### RNG / Initialization + +| megatron-lm arguments | Megatron Bridge config | Description | +| --- | --- | --- | +| `--seed` | `rng.seed` | Global random seed. | +| `--data-parallel-random-init` | `rng.data_parallel_random_init` | Enable per-DP-rank random init. | +| `--te-rng-tracker` | `rng.te_rng_tracker` | Use TE RNG (needed for CUDA graphs). | +| `--inference-rng-tracker` | `rng.inference_rng_tracker` | RNG tuned for inference stability. | + +### Distributed init and topology + +| megatron-lm arguments | Megatron Bridge config | Description | +| --- | --- | --- | +| `--distributed-backend` | `dist.distributed_backend` | Process group backend (nccl/gloo). | +| `--distributed-timeout-minutes` | `dist.distributed_timeout_minutes` | PG init and collective timeout. | +| `--no-align-grad-reduce` | `dist.align_grad_reduce=false` | Launch DP reduces independently per PP stage. | +| `--disable-gloo-process-groups` | `dist.use_gloo_process_groups=false` | Disable auxiliary Gloo PG creation. | +| `--use-sharp` | `dist.use_sharp` | Enable SHARP collectives for DP PG. | +| `--sharp-enabled-group` | `dist.sharp_enabled_group` | Which DP group enables SHARP. | +| `--high-priority-stream-groups` | `dist.high_priority_stream_groups` | Use high-priority comm streams for groups. | +| `--use-tp-pp-dp-mapping` | `dist.use_tp_pp_dp_mapping` | Use TP-PP-DP rank ordering at init. | + +Additional distributed/optimizer overlap settings: + +| megatron-lm arguments | Megatron Bridge config | Description | +| --- | --- | --- | +| `--use-distributed-optimizer` | `ddp.use_distributed_optimizer` and `optimizer.use_distributed_optimizer` | Enable distributed optimizer; settings are synchronized. | +| `--overlap-grad-reduce` | `ddp.overlap_grad_reduce` | Overlap DP gradient reduce-scatter. | +| `--overlap-param-gather` | `ddp.overlap_param_gather` | Overlap parameter all-gather with fprop. | + +### Profiling + +| megatron-lm arguments | Megatron Bridge config | Description | +| --- | --- | --- | +| `--profile` | `profiling.use_nsys_profiler` | Enable nsys profiling (capture is controlled via external CLI). | +| `--use-pytorch-profiler` | `profiling.use_pytorch_profiler` | Enable PyTorch profiler (TB-friendly). | +| `--profile-step-start` | `profiling.profile_step_start` | Global step to start profiling. | +| `--profile-step-end` | `profiling.profile_step_end` | Global step to stop profiling. | +| `--profile-ranks` | `profiling.profile_ranks` | Global ranks to profile. | +| `--record-memory-history` | `profiling.record_memory_history` | Track memory history. | +| `--memory-snapshot-path` | `profiling.memory_snapshot_path` | Output path for memory snapshot. | +| (shapes) | `profiling.record_shapes` | Record tensor shapes (overhead). | + +### In-process restart + +| megatron-lm arguments | Megatron Bridge config | Description | +| --- | --- | --- | +| `--inprocess-restart` | `inprocess_restart.enabled` | Enable nvrx in-process restart. | +| `--inprocess-max-iterations` | `inprocess_restart.max_iterations` | Max restart attempts. | +| `--inprocess-monitor-thread-interval` | `inprocess_restart.monitor_thread_interval` | Monitor thread polling interval. | +| `--inprocess-monitor-process-interval` | `inprocess_restart.monitor_process_interval` | Monitor process polling interval. | +| `--inprocess-progress-watchdog-interval` | `inprocess_restart.progress_watchdog_interval` | Auto progress timestamp update cadence. | +| `--inprocess-heartbeat-interval` | `inprocess_restart.heartbeat_interval` | Unresponsive-rank heartbeat cadence. | +| `--inprocess-soft-timeout` | `inprocess_restart.soft_timeout` | Soft progress timeout. | +| `--inprocess-hard-timeout` | `inprocess_restart.hard_timeout` | Hard timeout until kill. | +| `--inprocess-heartbeat-timeout` | `inprocess_restart.heartbeat_timeout` | Missing heartbeat timeout. | +| `--inprocess-barrier-timeout` | `inprocess_restart.barrier_timeout` | Timeout for internal barriers. | +| `--inprocess-completion-timeout` | `inprocess_restart.completion_timeout` | Timeout for completion barrier. | +| `--inprocess-last-call-wait` | `inprocess_restart.last_call_wait` | Delay to collect terminal failures. | +| `--inprocess-termination-grace-time` | `inprocess_restart.termination_grace_time` | SIGTERM→SIGKILL grace period. | +| `--inprocess-granularity` | `inprocess_restart.granularity` | Restart granularity (node/rank). | +| `--inprocess-active-world-size` | `inprocess_restart.active_world_size` | Active ranks count; rest are reserve. | +| `--inprocess-empty-cuda-cache` | `inprocess_restart.empty_cuda_cache` | Empty CUDA cache on restart finalize. | + +### Straggler detection + +| megatron-lm arguments | Megatron Bridge config | Description | +| --- | --- | --- | +| `--log-straggler` | `straggler.log_straggler` | Track and log straggler GPUs. | +| `--disable-straggler-on-startup` | `straggler.disable_straggler_on_startup` | Start with straggler detector disabled. | +| `--straggler-ctrlr-port` | `straggler.straggler_ctrlr_port` | Controller port for toggling. | +| `--straggler-minmax-count` | `straggler.straggler_minmax_count` | Num ranks to report for min/max throughput. | + +### Rerun state machine + +| megatron-lm arguments | Megatron Bridge config | Description | +| --- | --- | --- | +| `--error-injection-rate` | `rerun_state_machine.error_injection_rate` | Frequency of injected validation perturbations. | +| `--error-injection-type` | `rerun_state_machine.error_injection_type` | Kind of injection (correct/transient/persistent). | +| `--rerun-mode` | `rerun_state_machine.rerun_mode` | Disabled/validate_results/report_determinism_stats. | + +### Data / Tokenizer args + +| megatron-lm arguments | Megatron Bridge config | Description | +| --- | --- | --- | +| `--tokenizer-type` | `tokenizer.tokenizer_type` | Tokenizer implementation (e.g., HuggingFaceTokenizer). | +| `--tokenizer-model` | `tokenizer.tokenizer_model` | Model name/path for tokenizer. | +| `--num-workers` | `dataset.num_workers` | DataLoader workers. | +| `--no-create-attention-mask-in-dataloader` | `dataset.skip_getting_attention_mask_from_dataset=true` | Use backend-generated masks. | diff --git a/fern/v0.2.0/pages/modelopt/README.mdx b/fern/v0.2.0/pages/modelopt/README.mdx new file mode 100644 index 0000000000..4e73f30e27 --- /dev/null +++ b/fern/v0.2.0/pages/modelopt/README.mdx @@ -0,0 +1,28 @@ +--- +title: "Model Optimization" +description: "" +--- + +This directory contains comprehensive documentation for optimizing models with Megatron Bridge using NVIDIA ModelOpt. Learn how to apply various optimization techniques to improve inference efficiency while maintaining model quality. + +## Overview + +NVIDIA ModelOpt provides a suite of model optimization techniques for improving inference performance: + +- **Quantization** - Convert models from high-precision (FP32/BF16) to lower-precision formats (FP8, INT8, INT4) for efficient deployment +- **Distillation** - Transfer knowledge from a pre-trained teacher model to a smaller, faster student model +- **Pruning** - Reduce model size by removing layers (depth) or reducing dimensions (width) such as attention heads and hidden sizes + +## Quick Navigation + +### I want to + +**🔧 Quantize a pretrained model** +→ See the [Post-Training Quantization section](/quantization#post-training-quantization-ptq) for complete PTQ workflows (quantize, resume and generate, export) + +**🏋️ Train with quantization** +→ Check the [Quantization-Aware Training section](/quantization#quantization-aware-training-qat) for QAT workflows + +## References + +- [NVIDIA ModelOpt](https://github.com/NVIDIA/TensorRT-Model-Optimizer) diff --git a/fern/v0.2.0/pages/modelopt/quantization.mdx b/fern/v0.2.0/pages/modelopt/quantization.mdx new file mode 100644 index 0000000000..950f1f9a99 --- /dev/null +++ b/fern/v0.2.0/pages/modelopt/quantization.mdx @@ -0,0 +1,216 @@ +--- +title: "Quantization" +description: "" +--- + +This guide covers model quantization in Megatron Bridge using NVIDIA ModelOpt, including post-training quantization (PTQ) and quantization-aware training (QAT). + +## Table of Contents + +- [Overview](#overview) +- [Post-Training Quantization (PTQ)](#post-training-quantization-ptq) +- [Quantization-Aware Training (QAT)](#quantization-aware-training-qat) + +## Overview + +Quantization is an effective model optimization technique that compresses models by reducing precision from high-precision formats (FP16/BF16) to lower-precision formats (NVFP4, FP8, INT8, INT4). Quantization with Model Optimizer can compress model size by 2x-4x, speeding up inference while preserving model quality. + +In Megatron Bridge, quantization is enabled by NVIDIA Model Optimizer (ModelOpt) — a library to quantize and compress deep learning models for optimized inference on GPUs. Model Optimizer enables highly performant quantization formats including FP8, INT8, INT4, and NVFP4, and supports advanced algorithms such as SmoothQuant and AWQ with easy-to-use Python APIs. + +### Quantization Methods + +Megatron Bridge supports two quantization approaches: + +#### Post-Training Quantization (PTQ) + +PTQ reduces model precision after training to improve inference efficiency without requiring retraining. This is the fastest approach and works well for most models. + +**Process:** +1. Load a pretrained model checkpoint +2. Calibrate the model using a small dataset (typically 128-512 samples) to obtain scaling factors +3. Produce a quantized checkpoint + +#### Quantization-Aware Training (QAT) + +Quantization Aware Training (QAT) helps to improve the model accuracy beyond post training quantization (PTQ). QAT can further preserve model accuracy at low precisions (e.g., INT4, or FP4 in NVIDIA Blackwell platform). + +**Process:** + +1. Train/fine-tune the model in the original precision without quantization +2. Quantize the model from step 1 with `mtq.quantize()` +3. Train/fine-tune the quantized model with a small learning rate (e.g., 1e-5 for Adam optimizer) + +> **Note**: Step 3 is the actual 'Quantization Aware Training' step. The optimal hyperparameter setting for QAT can vary depending on the model and training dataset. + +> **Note**: QAT without the original precision training/fine-tuning (i.e., skipping Step 1) gives worse accuracy. Therefore, un-quantized original precision training/fine-tuning followed by QAT is recommended for best accuracy. + +--- + +## Post-Training Quantization (PTQ) + +PTQ quantizes a pretrained model by running calibration with a small dataset to compute scaling factors. The complete workflow includes: quantize → resume and generate → export. + +### Quantize + +Use the `examples/quantization/quantize.py` script for LLM PTQ: + +```bash +torchrun --nproc_per_node 2 examples/quantization/quantize.py \ + --hf-model-id meta-llama/Llama-3.2-1B \ + --export-quant-cfg fp8 \ + --tp 2 \ + --megatron-save-path ./llama3_2_1b_fp8 +``` + +Use the `examples/quantization/quantize_vlm.py` script for VLM PTQ: + +```bash +torchrun --nproc_per_node 8 examples/quantization/quantize_vlm.py \ + --hf-model-id Qwen/Qwen3-VL-30B-A3B-Instruct \ + --export-quant-cfg fp8 \ + --megatron-save-path ./Qwen3-VL-30B-A3B-Instruct_fp8 \ + --tp 4 \ + --etp 4 \ + --pp 2 \ + --calib-size 256 +``` + +**Key Arguments:** +- `--hf-model-id` - HuggingFace model ID or local path +- `--export-quant-cfg` - Quantization format (`fp8`, `nvfp4`, etc.) +- `--megatron-save-path` - Output checkpoint path +- `--tp` - Tensor parallelism size +- `--pp` - Pipeline parallelism size +- `--ep` - Expert parallelism for MoE models +- `--etp` - Expert tensor parallelism for MoE models +- `--calib-size` - Calibration samples for quantization (default: 512) + +### Resume and Generate + +Resume the quantized checkpoint and test with text generation using `examples/quantization/ptq_generate.py` for LLM: + +```bash +torchrun --nproc_per_node 2 examples/quantization/ptq_generate.py \ + --hf-model-id meta-llama/Llama-3.2-1B \ + --megatron-load-path ./llama3_2_1b_fp8 \ + --tp 2 +``` + +Resume the quantized checkpoint and test with text generation using `examples/quantization/ptq_generate_vlm.py` for VLM: + +```bash +torchrun --nproc_per_node 8 examples/quantization/ptq_generate_vlm.py \ + --hf-model-id Qwen/Qwen3-VL-30B-A3B-Instruct \ + --megatron-load-path ./Qwen3-VL-30B-A3B-Instruct_fp8 \ + --tp 8 \ + --ep 8 \ + --image-path ./demo.jpeg \ + --prompts "Describe this image." +``` + +**Key Arguments:** +- `--megatron-load-path` - Path to quantized checkpoint +- `--hf-model-id` - HuggingFace model ID or local path (for tokenizer) +- `--image-path` - Path to the input image file used for visual-language model prompt generation +- `--prompts` - Test prompts + +### Export + +Export the quantized checkpoint to unified HuggingFace format using `examples/quantization/export.py`: + +```bash +torchrun --nproc_per_node 2 examples/quantization/export.py \ + --hf-model-id meta-llama/Llama-3.2-1B \ + --megatron-load-path ./llama3_2_1b_fp8 \ + --export-dir ./llama3_2_1b_fp8_hf \ + --pp 2 \ + --dtype bfloat16 +``` + +**Key Arguments:** +- `--export-dir` - Output directory for unified HuggingFace checkpoint +- `--dtype` - Export data type + +### Supported Models For PTQ + +| Model | fp8 | nvfp4 | +|-------|-----|-------| +| Llama-3.2-1B | ✅ | ✅ | +| Qwen3-8B | ✅ | ✅ | +| Qwen3-30B-A3B | ✅ | ✅ | +| Nemotron-H-8B-Base-8K | ✅ | ✅ | +| Qwen3-VL-8B-Instruct | ✅ | ✅ | +| Qwen3-VL-30B-A3B-Instruct | ✅ | ✅ | + +--- + +## Quantization-Aware Training (QAT) + +In QAT, a model quantized using `mtq.quantize()` can be directly fine-tuned with the original training pipeline. During QAT, the scaling factors inside quantizers are frozen and the model weights are fine-tuned. + +### Complete QAT Workflow + +#### Step 1: Create Initial Quantized Checkpoint (PTQ) + +```bash +torchrun --nproc_per_node 8 examples/quantization/quantize.py \ + --hf-model-id meta-llama/Meta-Llama-3-8B \ + --export-quant-cfg fp8 \ + --tp 8 \ + --megatron-save-path /models/llama3_8b_fp8_init +``` + +#### Step 2: Configure Training + +Create a YAML configuration file (e.g., `conf/my_qat_config.yaml`): + +```yaml +model: + tensor_model_parallel_size: 4 + gradient_accumulation_fusion: False + +train: + train_iters: 20 + global_batch_size: 8 + eval_iters: 0 + +scheduler: + lr_warmup_iters: 10 + +logger: + log_interval: 1 + +checkpoint: + pretrained_checkpoint: /models/llama3_8b_fp8_init + save_interval: 20 + finetune: true +``` + +#### Step 3: Run QAT Training + +Use `examples/quantization/pretrain_quantized_llama3_8b.py`: + +```bash +python pretrain_quantized_llama3_8b.py \ + --nproc-per-node=4 \ + --config-file=conf/my_qat_config.yaml \ + --hf-path=meta-llama/Meta-Llama-3-8B +``` + +**Configuration Overrides:** + +You can also use command-line overrides: + +```bash +torchrun pretrain_quantized_llama3_8b.py \ + --nproc_per_node 4 \ + model.tensor_model_parallel_size=4 \ + model.gradient_accumulation_fusion=False \ + checkpoint.pretrained_checkpoint=/models/llama3_8b_fp8_init +``` + +### Supported Models For QAT + +| Model | Support | +|-------|---------| +| Meta-Llama-3-8B | ✅ | diff --git a/fern/v0.2.0/pages/models/README.mdx b/fern/v0.2.0/pages/models/README.mdx new file mode 100644 index 0000000000..0fd2aff8e5 --- /dev/null +++ b/fern/v0.2.0/pages/models/README.mdx @@ -0,0 +1,242 @@ +--- +title: "Supported Models" +description: "" +--- + +This directory contains documentation for all models supported by Megatron Bridge, including Large Language Models (LLMs) and Vision Language Models (VLMs). Each model documentation includes architecture details, conversion examples for Hugging Face ↔ Megatron Bridge, and links to training recipes. + +## Model Categories + +Megatron Bridge supports two main categories of models: + +### 🔤 Large Language Models (LLMs) + +Text-only models for language understanding and generation tasks. + +| Category | Model Count | Documentation | +|----------|-------------|---------------| +| **Large Language Models** | 13 models | [LLM Documentation](/llm/README) | + +**Supported LLM Families:** + +- DeepSeek (V2, V3) +- Gemma (2, 3) +- GLM-4.5 +- GPT-OSS +- LLaMA (3, Nemotron) +- Mistral +- Moonlight +- Nemotron-H +- OLMoE +- Qwen (2, 2.5, 3, 3 MoE, 3-Next) + +### 🖼️ Vision Language Models (VLMs) + +Multimodal models that combine vision and language capabilities. + +| Category | Model Count | Documentation | +|----------|-------------|---------------| +| **Vision Language Models** | 4 models | [VLM Documentation](/vlm/README) | + +**Supported VLM Families:** + +- Gemma 3 VL +- Nemotron Nano V2 VL +- Qwen (2.5 VL, 3 VL) + +--- + +## Quick Navigation + +### I want to + +**🔍 Find a specific LLM model** +→ Browse [Large Language Models](/llm/README) documentation + +**🖼️ Find a specific VLM model** +→ Browse [Vision Language Models](/vlm/README) documentation + +**🔄 Convert models between formats** +→ See [Bridge Guide](/../bridge-guide) for Hugging Face ↔ Megatron conversion + +**🚀 Get started with training** +→ See [Training Documentation](/../training/README) for training guides + +**📚 Understand model architectures** +→ Each model page documents architecture-specific features and configurations + +**🔧 Add support for a new model** +→ Refer to [Adding New Models](/../adding-new-models) + +**📊 Use training recipes** +→ Read [Recipe Usage](/../recipe-usage) for pre-configured training recipes + +--- + +## Model Documentation Structure + +Each model documentation page typically includes: + +1. **Model Overview** - Architecture and key features +2. **Available Variants** - Supported model sizes and configurations +3. **Conversion Examples** - Converting between Hugging Face and Megatron formats +4. **Training Recipes** - Links to training configurations and examples +5. **Architecture Details** - Model-specific features and configurations + +--- + +## Common Tasks by Model Type + +### For LLM Models + +**Training:** + +- Pretraining on large corpora +- Supervised fine-tuning (SFT) +- Parameter-efficient fine-tuning (PEFT/LoRA) +- Preference optimization (DPO) + +**Deployment:** + +- Export to Hugging Face format +- Integration with inference engines +- Model serving and deployment + +**Use Cases:** + +- Text generation +- Question answering +- Conversational AI +- Code generation + +### For VLM Models + +**Training:** + +- Multimodal pretraining +- Vision-language alignment +- Fine-tuning on visual tasks + +**Deployment:** + +- Export to Hugging Face format +- Multimodal inference + +**Use Cases:** + +- Image captioning +- Visual question answering +- Document understanding +- Multimodal reasoning + +--- + +## Related Documentation + +### Getting Started + +- **[Main Documentation](/../README)** - Return to main documentation +- **[Bridge Guide](/../bridge-guide)** - Hugging Face ↔ Megatron conversion +- **[Bridge Tech Details](/../bridge-tech-details)** - Technical details of the bridge system + +### Training Resources + +- **[Training Documentation](/../training/README)** - Comprehensive training guides +- **[Configuration Container](/../training/config-container-overview)** - Training configuration +- **[Parallelisms Guide](/../parallelisms)** - Data and model parallelism strategies +- **[Performance Guide](/../performance-guide)** - Performance optimization + +### Advanced Topics + +- **[Adding New Models](/../adding-new-models)** - Extending model support +- **[Recipe Usage](/../recipe-usage)** - Using pre-configured training recipes +- **[Bridge RL Integration](/../bridge-rl-integration)** - Reinforcement learning integration +- **[PEFT](/../training/peft)** - Parameter-efficient fine-tuning + +--- + +## Model Support Overview + +### By Architecture Type + +**Decoder-Only (Autoregressive):** + +- GPT-style models (GPT-OSS) +- LLaMA family (LLaMA 3, LLaMA Nemotron) +- Qwen family (Qwen 2, 2.5, 3, 3-Next) +- Gemma family (Gemma 2, 3) +- DeepSeek family (DeepSeek V2, V3) +- Mistral, Moonlight, Nemotron-H, GLM-4.5 + +**Mixture-of-Experts (MoE):** + +- Qwen 3 MoE, Qwen 3-Next +- DeepSeek V2, V3 +- OLMoE + +**Vision-Language (Multimodal):** + +- Gemma 3 VL +- Qwen 2.5 VL, Qwen 3 VL +- Nemotron Nano V2 VL + +### By Provider + +**Meta/LLaMA:** + +- LLaMA 3 + +**NVIDIA:** + +- LLaMA Nemotron +- Nemotron-H +- Nemotron Nano V2 VL + +**Alibaba Cloud:** + +- Qwen (2, 2.5, 3, 3 MoE, 3-Next) +- Qwen VL (2.5, 3) + +**Google:** + +- Gemma (2, 3) +- Gemma 3 VL + +**DeepSeek:** + +- DeepSeek (V2, V3) + +**Other:** + +- Mistral AI (Mistral) +- GLM-4.5 +- GPT-OSS +- Moonlight +- OLMoE + +--- + +## Conversion Support + +All models support bidirectional conversion: + +- **Hugging Face → Megatron Bridge**: Load pretrained weights for training +- **Megatron Bridge → Hugging Face**: Export trained models for deployment + +Conversion features: + +- Automatic architecture detection +- Parallelism-aware conversion (TP/PP/VPP/CP/EP) +- Streaming and memory-efficient transfers +- Verification mechanisms for conversion accuracy + +Refer to the [Bridge Guide](/../bridge-guide) for detailed conversion instructions. + +--- + +**Ready to explore?** Choose a model category: + +- [Large Language Models (LLMs)](/llm/README) +- [Vision Language Models (VLMs)](/vlm/README) + +Or return to the [main documentation](/../README). diff --git a/fern/v0.2.0/pages/models/llm/README.mdx b/fern/v0.2.0/pages/models/llm/README.mdx new file mode 100644 index 0000000000..9a5f87052f --- /dev/null +++ b/fern/v0.2.0/pages/models/llm/README.mdx @@ -0,0 +1,68 @@ +--- +title: "Large Language Models" +description: "" +--- + +This directory contains documentation for Large Language Models (LLMs) supported by Megatron Bridge. Each model documentation includes examples for converting to/from 🤗 Hugging Face and links to training recipes. + +## Available Models + +Megatron Bridge supports the following LLM families: + +| Model | Documentation | Description | +|-------|---------------|-------------| +| **DeepSeek V2** | [deepseek-v2.md](/deepseek-v2) | DeepSeek V2 model family | +| **DeepSeek V3** | [deepseek-v3.md](/deepseek-v3) | DeepSeek V3 model family | +| **Gemma 2** | [gemma2.md](/gemma2) | Google Gemma 2 models | +| **Gemma 3** | [gemma3.md](/gemma3) | Google Gemma 3 models | +| **GLM-4.5** | [glm45.md](/glm45) | GLM-4.5 model family | +| **GPT-OSS** | [gpt-oss.md](/gpt-oss) | Open-source GPT-style models | +| **LLaMA 3** | [llama3.md](/llama3) | Meta LLaMA 3 models | +| **LLaMA Nemotron** | [llama-nemotron.md](/llama-nemotron) | NVIDIA LLaMA Nemotron models | +| **Mistral** | [mistral.md](/mistral) | Mistral AI models | +| **Moonlight** | [moonlight.md](/moonlight) | Moonlight model family | +| **Nemotron-3** | [nemotron3.md](/nemotron3) | NVIDIA Nemotron-3 models | +| **Nemotron-H** | [nemotronh.md](/nemotronh) | NVIDIA Nemotron-H models | +| **OLMoE** | [olmoe.md](/olmoe) | OLMoE (Open Language Model - Mixture of Experts) | +| **Qwen** | [qwen.md](/qwen) | Alibaba Cloud Qwen model family | + +## Quick Navigation + +### I want to + +**🔍 Find a specific model** +→ Browse the model list above or use the [index page](/index) + +**🔄 Convert models between formats** +→ Each model page includes conversion examples for Hugging Face ↔ Megatron Bridge + +**🚀 Get started with training** +→ See [Training Documentation](/../../training/README) for training guides + +**📚 Understand model architecture** +→ Each model page documents architecture-specific features and configurations + +**🔧 Add support for a new model** +→ Refer to [Adding New Models](/../../adding-new-models) + +## Related Documentation + +- **[Models Overview](/../README)** - Return to main models documentation +- **[Vision Language Models](/../vlm/README)** - VLM model documentation +- **[Training Documentation](/../../training/README)** - Training and customization guides +- **[Bridge Guide](/../../bridge-guide)** - Working with Hugging Face models +- **[Adding New Models](/../../adding-new-models)** - Extending model support + +## Model Documentation Structure + +Each model documentation page typically includes: + +1. **Model Overview** - Architecture and key features +2. **Available Variants** - Supported model sizes and configurations +3. **Conversion Examples** - Converting between Hugging Face and Megatron formats +4. **Training Recipes** - Links to training configurations and examples +5. **Architecture Details** - Model-specific features and configurations + +--- + +**Ready to explore?** Choose a model from the list above or return to the [main documentation](/../../README). diff --git a/fern/v0.2.0/pages/models/llm/deepseek-v2.mdx b/fern/v0.2.0/pages/models/llm/deepseek-v2.mdx new file mode 100644 index 0000000000..0f09e17838 --- /dev/null +++ b/fern/v0.2.0/pages/models/llm/deepseek-v2.mdx @@ -0,0 +1,162 @@ +--- +title: "DeepSeek V2" +description: "" +--- + +[DeepSeek-V2](https://github.com/deepseek-ai/DeepSeek-V2) is a Mixture-of-Experts (MoE) language model that uses innovative Multi-head Latent Attention (MLA) for efficient inference and DeepSeekMoE architecture for economical training and inference. The model achieves performance comparable to GPT-4 while using significantly fewer activated parameters. More information is available in the companion paper ["DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model"](https://arxiv.org/abs/2405.04434). + +DeepSeek V2 models are supported via the Bridge system with auto-detected configuration and weight mapping. + +## Available Models + +Megatron Bridge supports the following DeepSeek V2 model variants: + +- **DeepSeek-V2**: 236B parameters (21B activated per token) +- **DeepSeek-V2-Lite**: 16B parameters (2.4B activated per token) + +Both models support pretraining with expert parallelism for efficient MoE training. + +## Model Architecture Features + +- **Multi-head Latent Attention (MLA)**: Novel attention mechanism that reduces KV cache requirements +- **DeepSeekMoE**: Efficient MoE architecture with routed and shared experts +- **Expert Parallelism**: Distributes experts across GPUs for scalable training +- **RoPE Embeddings**: Rotary position embeddings for position encoding +- **128K Context Length**: Native support for long sequences (DeepSeek-V2) +- **Pre-normalization**: RMSNorm before each transformer sub-layer + +## Conversion with 🤗 Hugging Face + +### Load HF → Megatron + +```python +from megatron.bridge import AutoBridge + +# Example: DeepSeek-V2-Lite +bridge = AutoBridge.from_hf_pretrained("deepseek-ai/DeepSeek-V2-Lite", trust_remote_code=True) +provider = bridge.to_megatron_provider() + +# Optionally configure parallelism before instantiating the model +provider.tensor_model_parallel_size = 1 +provider.pipeline_model_parallel_size = 1 +provider.expert_model_parallel_size = 8 + +model = provider.provide_distributed_model(wrap_with_ddp=False) +``` + +### Import Checkpoint from HF + +```bash +python examples/conversion/convert_checkpoints.py import \ + --hf-model deepseek-ai/DeepSeek-V2-Lite \ + --megatron-path /checkpoints/deepseek_v2_lite_megatron \ + --trust-remote-code +``` + +### Export Megatron → HF + +```python +from megatron.bridge import AutoBridge + +# Load the bridge from HF model ID +bridge = AutoBridge.from_hf_pretrained("deepseek-ai/DeepSeek-V2-Lite", trust_remote_code=True) + +# Export a trained Megatron checkpoint to HF format +bridge.export_ckpt( + megatron_path="/results/deepseek_v2_lite/checkpoints/iter_0000500", + hf_path="/exports/deepseek_v2_lite_hf", +) +``` + +### Run Inference on Converted Checkpoint + +```bash +python examples/conversion/hf_to_megatron_generate_text.py \ + --hf_model_path deepseek-ai/DeepSeek-V2-Lite \ + --megatron_model_path /checkpoints/deepseek_v2_lite_megatron \ + --prompt "What is artificial intelligence?" \ + --max_new_tokens 100 \ + --ep 8 \ + --trust-remote-code +``` + +For more details, see [examples/conversion/hf_to_megatron_generate_text.py](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/conversion/hf_to_megatron_generate_text.py) + +## Recipes + +See: [bridge.recipes.deepseek.deepseek_v2](/../../apidocs/bridge/bridge.recipes.deepseek.deepseek_v2) + +### Available Recipes + +- **Pretrain recipes**: + - `deepseek_v2_lite_pretrain_config`: Pre-training for DeepSeek-V2-Lite (16B parameters, 2.4B activated per token) + - `deepseek_v2_pretrain_config`: Pre-training for DeepSeek-V2 (236B parameters, 21B activated per token) + +### Parallelism Configurations + +| Model | TP | PP | EP | Total GPUs | Use Case | +|-------|----|----|----|-----------:|----------| +| **DeepSeek-V2-Lite** | 1 | 1 | 8 | 8 | Pre-training (single node) | +| **DeepSeek-V2** | 1 | 4 | 32 | 128 | Pre-training (16 nodes) | + +**Key Features**: +- **Expert Parallelism**: EP=8 (V2-Lite) or EP=32 (V2) for efficient MoE training +- **Selective Recomputation**: Enabled by default for memory optimization +- **Sequence Length**: Default 4096, V2 supports up to 128K tokens + +### Pre-training Examples + +#### DeepSeek-V2-Lite (16B) + +```python +from megatron.bridge.recipes.deepseek import deepseek_v2_lite_pretrain_config + +config = deepseek_v2_lite_pretrain_config( + name="deepseek_v2_lite_pretrain", + data_paths=["/path/to/dataset.nvjsonl"], + dir="/results/deepseek_v2_lite", + train_iters=500_000, + global_batch_size=512, + seq_length=4096, + # Uses TP=1, PP=1, EP=8 (8 GPUs) automatically +) +``` + +#### DeepSeek-V2 (236B) + +```python +from megatron.bridge.recipes.deepseek import deepseek_v2_pretrain_config + +config = deepseek_v2_pretrain_config( + name="deepseek_v2_pretrain", + data_paths=["/path/to/dataset.nvjsonl"], + dir="/results/deepseek_v2", + train_iters=500_000, + global_batch_size=512, + seq_length=4096, + # Uses TP=1, PP=4, EP=32 (128 GPUs) automatically +) +``` + +### Finetuning Recipes + +Finetuning recipes for DeepSeek V2 models are not currently available. + +## Hugging Face Model Cards & References + +### Hugging Face Model Cards +- DeepSeek-V2: https://huggingface.co/deepseek-ai/DeepSeek-V2 +- DeepSeek-V2-Lite: https://huggingface.co/deepseek-ai/DeepSeek-V2-Lite +- DeepSeek-V2-Chat: https://huggingface.co/deepseek-ai/DeepSeek-V2-Chat +- DeepSeek-V2-Lite-Chat: https://huggingface.co/deepseek-ai/DeepSeek-V2-Lite-Chat + +### Technical Papers +- DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model: [arXiv:2405.04434](https://arxiv.org/abs/2405.04434) + +### Additional Resources +- GitHub Repository: https://github.com/deepseek-ai/DeepSeek-V2 + +## Related Docs +- Recipe usage: [Recipe usage](/../../recipe-usage) +- Customizing the training recipe configuration: [Configuration overview](/../../training/config-container-overview) +- Training entry points: [Entry points](/../../training/entry-points) diff --git a/fern/v0.2.0/pages/models/llm/deepseek-v3.mdx b/fern/v0.2.0/pages/models/llm/deepseek-v3.mdx new file mode 100644 index 0000000000..a71c83ba4d --- /dev/null +++ b/fern/v0.2.0/pages/models/llm/deepseek-v3.mdx @@ -0,0 +1,154 @@ +--- +title: "DeepSeek V3" +description: "" +--- + +[DeepSeek-V3](https://github.com/deepseek-ai/DeepSeek-V3) is a large-scale Mixture-of-Experts (MoE) language model with 671B total parameters and 37B activated parameters per token. It features Multi-head Latent Attention (MLA), innovative load balancing strategies, and Multi-Token Prediction (MTP) for improved training efficiency. DeepSeek-V3 achieves state-of-the-art performance while maintaining economical training costs. More information is available in the technical report ["DeepSeek-V3 Technical Report"](https://arxiv.org/abs/2412.19437). + +DeepSeek V3 models are supported via the Bridge system with auto-detected configuration and weight mapping. + +## Available Models + +Megatron Bridge supports the following DeepSeek V3 model variants: + +- **DeepSeek-V3**: 671B parameters (37B activated per token) +- **DeepSeek-V3-Base**: Pre-trained base model without instruction tuning + +The model supports pretraining with expert parallelism, pipeline parallelism, and optional Multi-Token Prediction (MTP). + +## Model Architecture Features + +- **Multi-head Latent Attention (MLA)**: Advanced attention mechanism for reduced KV cache and improved efficiency +- **DeepSeekMoE**: Enhanced MoE architecture with 256 routed experts and shared experts +- **Multi-Token Prediction (MTP)**: Auxiliary training objective that predicts multiple future tokens +- **Expert Parallelism**: Distributes 256 experts across GPUs for scalable training +- **RoPE Embeddings**: Rotary position embeddings with scaling factor for position encoding +- **Sigmoid Gating with Expert Bias**: Novel routing mechanism with learnable expert bias +- **Pre-normalization**: RMSNorm before each transformer sub-layer for training stability + +## Conversion with 🤗 Hugging Face + +### Load HF → Megatron + +```python +from megatron.bridge import AutoBridge + +# Example: DeepSeek-V3-Base +bridge = AutoBridge.from_hf_pretrained("deepseek-ai/DeepSeek-V3-Base", trust_remote_code=True) +provider = bridge.to_megatron_provider() + +# Optionally configure parallelism before instantiating the model +provider.tensor_model_parallel_size = 2 +provider.pipeline_model_parallel_size = 16 +provider.expert_model_parallel_size = 64 + +model = provider.provide_distributed_model(wrap_with_ddp=False) +``` + +### Import Checkpoint from HF + +```bash +python examples/conversion/convert_checkpoints.py import \ + --hf-model deepseek-ai/DeepSeek-V3-Base \ + --megatron-path /checkpoints/deepseek_v3_megatron \ + --trust-remote-code +``` + +### Export Megatron → HF + +```python +from megatron.bridge import AutoBridge + +# Load the bridge from HF model ID +bridge = AutoBridge.from_hf_pretrained("deepseek-ai/DeepSeek-V3-Base", trust_remote_code=True) + +# Export a trained Megatron checkpoint to HF format +bridge.export_ckpt( + megatron_path="/results/deepseek_v3/checkpoints/iter_0000500", + hf_path="/exports/deepseek_v3_hf", +) +``` + +### Run Inference on Converted Checkpoint + +```bash +python examples/conversion/hf_to_megatron_generate_text.py \ + --hf_model_path deepseek-ai/DeepSeek-V3-Base \ + --megatron_model_path /checkpoints/deepseek_v3_megatron \ + --prompt "What is artificial intelligence?" \ + --max_new_tokens 100 \ + --tp 2 \ + --pp 16 \ + --ep 64 \ + --trust-remote-code +``` + +For more details, see [examples/conversion/hf_to_megatron_generate_text.py](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/conversion/hf_to_megatron_generate_text.py) + +## Recipes + +See: [bridge.recipes.deepseek.deepseek_v3](/../../apidocs/bridge/bridge.recipes.deepseek.deepseek_v3) + +### Available Recipes + +- **Pretrain recipes**: + - `deepseek_v3_pretrain_config`: Pre-training for DeepSeek-V3 (671B parameters, 37B activated per token) + +### Parallelism Configurations + +| Model | TP | PP | EP | VP | Nodes | Total GPUs | Use Case | +|-------|----|----|-----|-----|------:|-----------:|----------| +| **DeepSeek-V3** | 2 | 16 | 64 | None | 128 | 1024 | Pre-training | + +**Key Features**: +- **Expert Parallelism**: EP=64 for distributing 256 experts across GPUs +- **Pipeline Parallelism**: PP=16 with asymmetric layouts optimized for embedding and loss layers +- **Selective Recomputation**: Enabled by default for memory optimization +- **Multi-Token Prediction (MTP)**: Optional auxiliary training objective (1 layer by default) +- **Sequence Parallel**: Enabled by default for memory efficiency + +**Performance Optimizations**: +- **MoE Permute Fusion**: Fused expert permutation operations +- **Flex Dispatcher Backend**: Optional high-performance MoE token dispatcher +- **RoPE Fusion**: Optional fusion for Multi-head Latent Attention +- **Precision-Aware Optimizer**: FP32 master weights with BF16 gradients and optimizer states + +### Pre-training Example + +```python +from megatron.bridge.recipes.deepseek import deepseek_v3_pretrain_config + +config = deepseek_v3_pretrain_config( + name="deepseek_v3_pretrain", + data_paths=["/path/to/dataset.nvjsonl"], + dir="/results/deepseek_v3", + train_iters=500_000, + global_batch_size=4096, + seq_length=4096, + # MTP configuration + mtp_num_layers=1, + mtp_loss_scaling_factor=0.1, + # Uses TP=2, PP=16, EP=64 (1024 GPUs, 128 nodes) automatically +) +``` + +### Finetuning Recipes + +Finetuning recipes for DeepSeek V3 are not currently available. + +## Hugging Face Model Cards & References + +### Hugging Face Model Cards +- DeepSeek-V3: https://huggingface.co/deepseek-ai/DeepSeek-V3 +- DeepSeek-V3-Base: https://huggingface.co/deepseek-ai/DeepSeek-V3-Base + +### Technical Papers +- DeepSeek-V3 Technical Report: [arXiv:2412.19437](https://arxiv.org/abs/2412.19437) + +### Additional Resources +- GitHub Repository: https://github.com/deepseek-ai/DeepSeek-V3 + +## Related Docs +- Recipe usage: [Recipe usage](/../../recipe-usage) +- Customizing the training recipe configuration: [Configuration overview](/../../training/config-container-overview) +- Training entry points: [Entry points](/../../training/entry-points) diff --git a/fern/v0.2.0/pages/models/llm/gemma2.mdx b/fern/v0.2.0/pages/models/llm/gemma2.mdx new file mode 100644 index 0000000000..ceba2e9646 --- /dev/null +++ b/fern/v0.2.0/pages/models/llm/gemma2.mdx @@ -0,0 +1,292 @@ +--- +title: "Gemma 2" +description: "" +--- + +[Google's Gemma 2](https://huggingface.co/collections/google/gemma-2-release-667d6600fd5220e7b967f315) is a family of lightweight, open models built on the same research and technology used to create Gemini models. The Gemma 2 architecture builds on the transformer decoder framework with enhancements including pre-normalization with RMSNorm, GeGLU activations, Rotary Positional Embeddings (RoPE), attention logit softcapping, and sliding window attention. + +Gemma 2 models are designed for a wide range of text generation tasks and are available in multiple sizes to suit different computational budgets. + +Gemma family models are supported via the Bridge system with auto-detected configuration and weight mapping. + +## Available Models + +### Text-Only Models +- **Gemma 2 2B** (`google/gemma-2-2b`): Compact 2B parameter model optimized for efficiency + - 26 layers, 2304 hidden size + - 8 attention heads, 4 query groups (GQA) + - Sequence length: 8,192 tokens + - Ideal for single-GPU deployment + +- **Gemma 2 9B** (`google/gemma-2-9b`): Medium-scale 9B parameter model balancing performance and efficiency + - 42 layers, 3584 hidden size + - 16 attention heads, 8 query groups (GQA) + - Sequence length: 8,192 tokens + - Recommended: 4-8 GPUs with tensor parallelism + +- **Gemma 2 27B** (`google/gemma-2-27b`): Large-scale 27B parameter model for maximum performance + - 46 layers, 4608 hidden size + - 32 attention heads, 16 query groups (GQA) + - Sequence length: 8,192 tokens + - Recommended: 8-16 GPUs with tensor and pipeline parallelism + +All models support a sequence length of 8,192 tokens and use sliding window attention for efficient context processing. + +## Model Architecture Features + +Gemma 2 introduces several architectural innovations: + +- **Sliding Window Attention**: Local attention with a 4,096 token window for efficient long-context processing +- **Attention Logit Softcapping**: Caps attention logits at 50.0 to prevent extreme values +- **Final Logit Softcapping**: Caps final output logits at 30.0 for stable generation +- **GeGLU Activation**: Uses gated linear units with GELU activation for improved performance +- **RMSNorm**: Layer normalization without mean centering for faster computation +- **Rotary Embeddings**: RoPE with base frequency 10,000 +- **Query Pre-Attention Scaling**: Scales queries by 224 before attention computation + +## Conversion with 🤗 Hugging Face + +### Load HF → Megatron +```python +from megatron.bridge import AutoBridge + +# Example: Gemma 2 9B +bridge = AutoBridge.from_hf_pretrained("google/gemma-2-9b") +provider = bridge.to_megatron_provider() + +# Configure parallelism before instantiating the model +provider.tensor_model_parallel_size = 8 +provider.pipeline_model_parallel_size = 1 + +provider.finalize() +model = provider.provide_distributed_model(wrap_with_ddp=False) +``` + +### Import HF → Megatron +To import the HF model to your desired Megatron path: +```bash +# Gemma 2 2B +python examples/conversion/convert_checkpoints.py import \ +--hf-model google/gemma-2-2b \ +--megatron-path /models/gemma-2-2b + +# Gemma 2 9B +python examples/conversion/convert_checkpoints.py import \ +--hf-model google/gemma-2-9b \ +--megatron-path /models/gemma-2-9b + +# Gemma 2 27B +python examples/conversion/convert_checkpoints.py import \ +--hf-model google/gemma-2-27b \ +--megatron-path /models/gemma-2-27b +``` + +### Export Megatron → HF +```bash +# Gemma 2 9B example +python examples/conversion/convert_checkpoints.py export \ +--hf-model google/gemma-2-9b \ +--megatron-path /results/gemma2_9b/checkpoints/iter_00001000 \ +--hf-path ./gemma2-9b-hf-export +``` + +### Run Inference on Converted Checkpoint + +```bash +python examples/conversion/hf_to_megatron_generate_text.py \ +--hf_model_path google/gemma-2-9b \ +--megatron_model_path /models/gemma-2-9b \ +--prompt "What is artificial intelligence?" \ +--max_new_tokens 100 +``` + +Note: +- `--megatron_model_path` is optional. If not specified, the script will convert the model and then run forward. + +## Pretrain and Finetune Recipes + +- See: [bridge.recipes.gemma](/../../apidocs/bridge/bridge.recipes.gemma) +- Available recipes: + - **Pretraining:** + - `gemma2_2b_pretrain_config`: Pre-training configuration for Gemma 2 2B + - `gemma2_9b_pretrain_config`: Pre-training configuration for Gemma 2 9B + - `gemma2_27b_pretrain_config`: Pre-training configuration for Gemma 2 27B + - **Finetuning:** + - `gemma2_2b_finetune_config`: Finetuning configuration for Gemma 2 2B with PEFT support (LoRA, DoRA) + - `gemma2_9b_finetune_config`: Finetuning configuration for Gemma 2 9B with PEFT support (LoRA, DoRA) + - `gemma2_27b_finetune_config`: Finetuning configuration for Gemma 2 27B with PEFT support (LoRA, DoRA) + +Before training, ensure the following environment variables are set: +1. `SAVE_DIR`: checkpoint and log saving directory +2. `HF_TOKEN`: to download models from HF Hub (if required) +3. `HF_HOME`: (optional) to avoid re-downloading models and datasets +4. `WANDB_API_KEY`: (optional) to enable WandB logging + +### Pretraining + +#### Gemma 2 2B +```python +from megatron.bridge.recipes.gemma import gemma2_2b_pretrain_config + +# Create a pretraining configuration +config = gemma2_2b_pretrain_config( + name="my_gemma2_2b_pretrain", + data_paths=["path/to/data"], + train_iters=100000, + global_batch_size=32, +) +``` + +#### Gemma 2 9B +```python +from megatron.bridge.recipes.gemma import gemma2_9b_pretrain_config + +config = gemma2_9b_pretrain_config( + name="my_gemma2_9b_pretrain", + data_paths=["path/to/data"], + train_iters=100000, + global_batch_size=32, +) +``` + +#### Gemma 2 27B +```python +from megatron.bridge.recipes.gemma import gemma2_27b_pretrain_config + +config = gemma2_27b_pretrain_config( + name="my_gemma2_27b_pretrain", + data_paths=["path/to/data"], + train_iters=100000, + global_batch_size=32, +) +``` + +### Full Finetuning + +#### Gemma 2 2B +```bash +torchrun --nproc-per-node=8 run/run_recipe.py \ +--pretrained-checkpoint /models/gemma-2-2b \ +--recipe gemma2_2b_finetune_config \ +--peft none \ +train.global_batch_size=64 \ +train.train_iters=1000 \ +checkpoint.save=$SAVE_DIR/gemma2_2b_finetune +``` + +Or programmatically: +```python +from megatron.bridge.recipes.gemma import gemma2_2b_finetune_config + +config = gemma2_2b_finetune_config( + name="gemma2_2b_full_finetune", + pretrained_checkpoint="/models/gemma-2-2b", + peft="none", + train_iters=1000, + global_batch_size=64, +) +``` + +#### Gemma 2 9B +```bash +torchrun --nproc-per-node=8 run/run_recipe.py \ +--pretrained-checkpoint /models/gemma-2-9b \ +--recipe gemma2_9b_finetune_config \ +--peft none \ +train.global_batch_size=64 \ +train.train_iters=1000 \ +checkpoint.save=$SAVE_DIR/gemma2_9b_finetune +``` + +#### Gemma 2 27B +```bash +torchrun --nproc-per-node=16 run/run_recipe.py \ +--pretrained-checkpoint /models/gemma-2-27b \ +--recipe gemma2_27b_finetune_config \ +--peft none \ +train.global_batch_size=64 \ +train.train_iters=1000 \ +checkpoint.save=$SAVE_DIR/gemma2_27b_finetune +``` + +### Parameter-Efficient Finetuning (PEFT) with LoRA + +#### Gemma 2 2B +```bash +torchrun --nproc-per-node=8 run/run_recipe.py \ +--pretrained-checkpoint /models/gemma-2-2b \ +--recipe gemma2_2b_finetune_config \ +--peft lora \ +train.global_batch_size=128 \ +checkpoint.save=$SAVE_DIR/gemma2_2b_lora +``` + +PEFT options: +- `--peft`: Set to `lora` for LoRA or `dora` for DoRA. Set to `none` for full finetuning. + +Or programmatically: +```python +from megatron.bridge.recipes.gemma import gemma2_2b_finetune_config + +# LoRA finetuning +config = gemma2_2b_finetune_config( + name="gemma2_2b_lora_finetune", + pretrained_checkpoint="/models/gemma-2-2b", + peft="lora", # or "dora" + train_iters=1000, + global_batch_size=128, +) +``` + +#### Gemma 2 9B LoRA +```python +from megatron.bridge.recipes.gemma import gemma2_9b_finetune_config + +config = gemma2_9b_finetune_config( + name="gemma2_9b_lora_finetune", + pretrained_checkpoint="/models/gemma-2-9b", + peft="lora", + train_iters=1000, + global_batch_size=128, +) +``` + +#### Gemma 2 27B LoRA +```python +from megatron.bridge.recipes.gemma import gemma2_27b_finetune_config + +config = gemma2_27b_finetune_config( + name="gemma2_27b_lora_finetune", + pretrained_checkpoint="/models/gemma-2-27b", + peft="lora", + train_iters=1000, + global_batch_size=128, +) +``` + +### Recommended Configurations + +| Model | Mode | TP | PP | Global Batch Size | Learning Rate | +|-------|------|----|----|-------------------|---------------| +| Gemma 2 2B | Full SFT | 1 | 1 | 64-128 | 5e-6 | +| Gemma 2 2B | LoRA/DoRA | 1 | 1 | 128-256 | 1e-4 | +| Gemma 2 9B | Full SFT | 4 | 1 | 64-128 | 5e-6 | +| Gemma 2 9B | LoRA/DoRA | 1 | 1 | 128-256 | 1e-4 | +| Gemma 2 27B | Full SFT | 8 | 2 | 64-128 | 5e-6 | +| Gemma 2 27B | LoRA/DoRA | 4 | 1 | 128-256 | 1e-4 | + +## Examples +- Checkpoint import/export: [examples/conversion/convert_checkpoints.py](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/conversion/convert_checkpoints.py) +- Generate text (HF→Megatron): [examples/conversion/hf_to_megatron_generate_text.py](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/conversion/hf_to_megatron_generate_text.py) + +## Hugging Face Model Cards + +- Gemma 2 2B: https://huggingface.co/google/gemma-2-2b +- Gemma 2 9B: https://huggingface.co/google/gemma-2-9b +- Gemma 2 27B: https://huggingface.co/google/gemma-2-27b +- Gemma 2 Collection: https://huggingface.co/collections/google/gemma-2-release-667d6600fd5220e7b967f315 + +## Related Docs +- Recipe usage: [Recipe usage](/../../recipe-usage) +- Customizing the training recipe configuration: [Configuration overview](/../../training/config-container-overview) +- Training entry points: [Entry points](/../../training/entry-points) diff --git a/fern/v0.2.0/pages/models/llm/gemma3.mdx b/fern/v0.2.0/pages/models/llm/gemma3.mdx new file mode 100644 index 0000000000..b89ced5a21 --- /dev/null +++ b/fern/v0.2.0/pages/models/llm/gemma3.mdx @@ -0,0 +1,189 @@ +--- +title: "Gemma 3" +description: "" +--- + +[Google's Gemma 3](https://huggingface.co/collections/google/gemma-3-release) is a family of lightweight, state-of-the-art open models built on the same research and technology used to create Gemini models. The Gemma 3 architecture builds on the transformer decoder framework with enhancements including pre-normalization with RMSNorm, GeGLU activations, Rotary Positional Embeddings (RoPE), and hybrid attention patterns (sliding window and global attention). + +Gemma 3 models are designed for a wide range of text generation tasks and are available in multiple sizes to suit different computational budgets. + +Gemma family models are supported via the Bridge system with auto-detected configuration and weight mapping. + +## Available Models + +### Text-Only Models +- **Gemma 3 1B** (`google/gemma-3-1b-it`): Compact 1B parameter model optimized for efficiency + - 26 layers, 1152 hidden size + - 8 attention heads, 2 query groups (GQA) + - Sequence length: 131,072 tokens + - Ideal for single-GPU deployment + +All models support a sequence length of 131,072 tokens and use hybrid attention patterns (sliding window + global). + +## Model Architecture Features + +Gemma 3 introduces several architectural innovations: + +- **Hybrid Attention Pattern**: Alternates between global and local sliding window attention for efficient long-context processing +- **GeGLU Activation**: Uses gated linear units with GELU activation for improved performance +- **RMSNorm**: Layer normalization without mean centering for faster computation +- **Rotary Embeddings**: Separate RoPE configurations for local and global attention layers + - Local attention: Uses sliding window with rotary base 10,000 + - Global attention: Extended rotary base for better long-range dependencies + +## Conversion with 🤗 Hugging Face + +### Load HF → Megatron +```python +from megatron.bridge import AutoBridge + +# Example: Gemma 3 1B +bridge = AutoBridge.from_hf_pretrained("google/gemma-3-1b-it") +provider = bridge.to_megatron_provider() + +# Configure parallelism before instantiating the model +provider.tensor_model_parallel_size = 1 +provider.pipeline_model_parallel_size = 1 + +provider.finalize() +model = provider.provide_distributed_model(wrap_with_ddp=False) +``` + +### Import HF → Megatron +To import the HF model to your desired Megatron path: +```bash +python examples/conversion/convert_checkpoints.py import \ +--hf-model google/gemma-3-1b-it \ +--megatron-path /models/gemma-3-1b-it +``` + +### Export Megatron → HF +```bash +python examples/conversion/convert_checkpoints.py export \ +--hf-model google/gemma-3-1b-it \ +--megatron-path /results/gemma3_1b/checkpoints/iter_00001000 \ +--hf-path ./gemma3-hf-export +``` + +### Run Inference on Converted Checkpoint + +```bash +python examples/conversion/hf_to_megatron_generate_text.py \ +--hf_model_path google/gemma-3-1b-it \ +--megatron_model_path /models/gemma-3-1b-it \ +--prompt "What is artificial intelligence?" \ +--max_new_tokens 100 +``` + +Note: +- `--megatron_model_path` is optional. If not specified, the script will convert the model and then run forward. + +## Recipes + +See: [bridge.recipes.gemma](/../../apidocs/bridge/bridge.recipes.gemma) + +### Available Recipes + +- **Pretrain recipes**: + - `gemma3_1b_pretrain_config`: Pre-training for Gemma 3 1B + +- **Finetune recipes**: + - `gemma3_1b_finetune_config`: Finetuning for Gemma 3 1B with PEFT support (LoRA, DoRA) + +**Before training**, ensure these environment variables are set: +- `SAVE_DIR`: checkpoint and log saving directory +- `HF_TOKEN`: to download models from HF Hub (if required) +- `HF_HOME`: (optional) to avoid re-downloading models and datasets +- `WANDB_API_KEY`: (optional) to enable WandB logging + +### Parallelism Configurations + +| Model | Mode | TP | PP | Total GPUs | Use Case | +|-------|------|----|----|------------|----------| +| **Gemma 3 1B** | Pretrain | 1 | 1 | 8 | Pre-training (single node) | +| **Gemma 3 1B** | Full SFT | 1 | 1 | 8 | Full supervised finetuning | +| **Gemma 3 1B** | LoRA/DoRA | 1 | 1 | 8 | PEFT finetuning (single node) | + +### Pre-training Example + +```python +from megatron.bridge.recipes.gemma import gemma3_1b_pretrain_config + +config = gemma3_1b_pretrain_config( + name="gemma3_1b_pretrain", + data_paths=["path/to/data"], + train_iters=100000, + global_batch_size=256, + # Uses TP=1, PP=1 (8 GPUs) automatically +) +``` + +### Finetuning Examples + +#### Full Finetuning + +```python +from megatron.bridge.recipes.gemma import gemma3_1b_finetune_config + +config = gemma3_1b_finetune_config( + name="gemma3_1b_full_finetune", + pretrained_checkpoint="/models/gemma-3-1b-it", + peft=None, + train_iters=1000, + global_batch_size=64, + finetune_lr=5e-6, + # Uses TP=1, PP=1 (8 GPUs) automatically +) +``` + +#### LoRA Finetuning + +```python +from megatron.bridge.recipes.gemma import gemma3_1b_finetune_config + +config = gemma3_1b_finetune_config( + name="gemma3_1b_lora_finetune", + pretrained_checkpoint="/models/gemma-3-1b-it", + peft="lora", # or "dora" + train_iters=1000, + global_batch_size=128, + finetune_lr=1e-4, + # Uses TP=1, PP=1 (8 GPUs) automatically +) +``` + +### Command-Line Training + +**Full Finetuning:** +```bash +torchrun --nproc-per-node=8 run/run_recipe.py \ + --pretrained-checkpoint /models/gemma-3-1b-it \ + --recipe gemma3_1b_finetune_config \ + train.global_batch_size=64 \ + train.train_iters=1000 \ + checkpoint.save=$SAVE_DIR/gemma3_1b_finetune +``` + +**LoRA Finetuning:** +```bash +torchrun --nproc-per-node=8 run/run_recipe.py \ + --pretrained-checkpoint /models/gemma-3-1b-it \ + --recipe gemma3_1b_finetune_config \ + --peft_scheme lora \ + train.global_batch_size=128 \ + checkpoint.save=$SAVE_DIR/gemma3_1b_lora +``` + +## Examples +- Checkpoint import/export: [examples/conversion/convert_checkpoints.py](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/conversion/convert_checkpoints.py) +- Generate text (HF→Megatron): [examples/conversion/hf_to_megatron_generate_text.py](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/conversion/hf_to_megatron_generate_text.py) + +## Hugging Face Model Cards + +- Gemma 3 1B: https://huggingface.co/google/gemma-3-1b-it + +## Related Docs +- Gemma3 Vision-Language Models: [Gemma 3 VL](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/models/vlm/gemma3_vl/README.md) +- Recipe usage: [Recipe usage](/../../recipe-usage) +- Customizing the training recipe configuration: [Configuration overview](/../../training/config-container-overview) +- Training entry points: [Entry points](/../../training/entry-points) diff --git a/fern/v0.2.0/pages/models/llm/glm45.mdx b/fern/v0.2.0/pages/models/llm/glm45.mdx new file mode 100644 index 0000000000..2af014fc11 --- /dev/null +++ b/fern/v0.2.0/pages/models/llm/glm45.mdx @@ -0,0 +1,295 @@ +--- +title: "GLM 4.5" +description: "" +--- + +[GLM 4.5](https://huggingface.co/zai-org/GLM-4.5) is a family of large-scale Mixture-of-Experts (MoE) language models developed by Zhipu AI. Built on the GLM (General Language Model) architecture, GLM 4.5 introduces advanced features including sparse MoE layers, shared expert mechanisms, and Multi-Token Prediction (MTP) for improved training efficiency and performance. + +GLM 4.5 models are designed for high-performance text generation and understanding tasks, offering excellent quality with efficient inference through sparse activation. The family includes two variants optimized for different deployment scenarios. + +GLM 4.5 models are supported via the Bridge system with auto-detected configuration and weight mapping. + +## Available Models + +### GLM 4.5 355B-A32B +**Full Model:** `zai-org/GLM-4.5` (355B total parameters, 32B active per token) +- **Architecture:** + - 92 transformer layers (first 3 dense, remaining 89 MoE) + - 5120 hidden size, 12288 FFN hidden size + - 96 attention heads, 8 query groups (GQA) + - 160 experts per MoE layer, top-8 routing with 2.5x scaling factor + - MoE FFN hidden size: 1536 + - Shared expert intermediate size: 1536 + - Includes QK LayerNorm for improved training stability +- **Context:** 131,072 tokens +- **Vocabulary:** 151,552 tokens +- **Optimizations:** + - Shared expert overlap for better load balancing + - Router sigmoid scoring with expert bias + - MTP (Multi-Token Prediction) with 1 layer, 0.3 scaling factor + +### GLM 4.5 Air 106B-A12B +**Full Model:** `zai-org/GLM-4.5-Air` (106B total parameters, 12B active per token) +- **Architecture:** + - 46 transformer layers (first layer dense, remaining 45 MoE) + - 4096 hidden size, 10944 FFN hidden size + - 96 attention heads, 8 query groups (GQA) + - 128 experts per MoE layer, top-8 routing with 1.0x scaling factor + - MoE FFN hidden size: 1408 + - Shared expert intermediate size: 1408 + - No QK LayerNorm +- **Context:** 131,072 tokens +- **Vocabulary:** 151,552 tokens +- **Optimizations:** + - Optimized for reduced memory footprint + - Suitable for mid-range GPU clusters + +Both models use RMSNorm, SiLU activation, gated linear units, and RoPE with 1M base frequency. + +## Model Architecture Features + +GLM 4.5 introduces several advanced architectural innovations: + +- **Mixture-of-Experts (MoE)**: Sparse activation with 160/128 experts, activating only 8 per token for efficient scaling +- **Shared Expert Mechanism**: Dedicated shared experts with overlap for improved load balancing and knowledge transfer +- **Multi-Token Prediction (MTP)**: Predicts multiple future tokens simultaneously for better training efficiency + - Configurable MTP layers (default: 1 layer) + - Adjustable loss scaling factor (default: 0.3 for early training, 0.1 for later stages) +- **RoPE with Extended Context**: Rotary embeddings with 1M base frequency for robust long-context modeling +- **Advanced Router Design**: + - Sigmoid score function for better expert selection + - Expert bias with configurable update rate + - Auxiliary loss for load balancing +- **Grouped Query Attention (GQA)**: 8 query groups for efficient attention computation +- **RMSNorm**: Fast layer normalization without mean centering + +## Conversion with 🤗 Hugging Face + +### Load HF → Megatron +```python +from megatron.bridge import AutoBridge + +# Example: GLM 4.5 Air 106B +bridge = AutoBridge.from_hf_pretrained("zai-org/GLM-4.5-Air") +provider = bridge.to_megatron_provider() + +# Configure parallelism before instantiating the model +# For Air 106B: TP=1, PP=4, EP=8 (32 GPUs) +provider.tensor_model_parallel_size = 1 +provider.pipeline_model_parallel_size = 4 +provider.expert_model_parallel_size = 8 +provider.sequence_parallel = True + +provider.finalize() +model = provider.provide_distributed_model(wrap_with_ddp=False) +``` + +### Import HF → Megatron +```bash +# Import GLM 4.5 Air model +python examples/conversion/convert_checkpoints.py import \ +--hf-model zai-org/GLM-4.5-Air \ +--megatron-path /models/glm45-air-106b + +# Import GLM 4.5 355B model +python examples/conversion/convert_checkpoints.py import \ +--hf-model zai-org/GLM-4.5 \ +--megatron-path /models/glm45-355b +``` + +### Export Megatron → HF +```bash +python examples/conversion/convert_checkpoints.py export \ +--hf-model zai-org/GLM-4.5-Air \ +--megatron-path /results/glm45_air/checkpoints/iter_00001000 \ +--hf-path ./glm45-air-hf-export +``` + +### Run Inference on Converted Checkpoint +```bash +python examples/conversion/hf_to_megatron_generate_text.py \ +--hf_model_path zai-org/GLM-4.5-Air \ +--megatron_model_path /models/glm45-air-106b \ +--prompt "Explain quantum computing in simple terms." \ +--max_new_tokens 200 +``` + +## Pretrain and Finetune Recipes + +### Available Recipes +- Pretrain recipes: + - `glm45_355b_pretrain_config`: Pre-training for GLM 4.5 355B + - `glm45_air_106b_pretrain_config`: Pre-training for GLM 4.5 Air 106B +- Finetune recipes: + - `glm45_355b_finetune_config`: Finetuning for GLM 4.5 355B with PEFT support + - `glm45_air_106b_finetune_config`: Finetuning for GLM 4.5 Air 106B with PEFT support + +### Parallelism Configurations + +#### GLM 4.5 355B-A32B +| Mode | TP | PP | EP | Total GPUs | Use Case | +|------|----|----|----|-----------:|----------| +| **Pretrain** | 2 | 8 | 16 | 256 | Full pre-training | +| **PEFT (LoRA/DoRA)** | 2 | 4 | 4 | 32 | Parameter-efficient finetuning | +| **Full SFT** | 2 | 8 | 16 | 256 | Full supervised finetuning | + +#### GLM 4.5 Air 106B-A12B +| Mode | TP | PP | EP | Total GPUs | Use Case | +|------|----|----|----|-----------:|----------| +| **Pretrain** | 1 | 4 | 8 | 32 | Full pre-training | +| **PEFT (LoRA/DoRA)** | 1 | 2 | 4 | 8 | Parameter-efficient finetuning (1 node!) | +| **Full SFT** | 1 | 4 | 8 | 32 | Full supervised finetuning | + +### Pre-training Example + +```python +from megatron.bridge.recipes.glm import glm45_air_106b_pretrain_config + +# Create a pre-training configuration +config = glm45_air_106b_pretrain_config( + name="glm45_air_pretrain", + data_paths=["path/to/data"], + train_iters=100000, + global_batch_size=2048, + micro_batch_size=1, + lr=1e-4, + min_lr=1e-5, + # MTP configuration + mtp_num_layers=1, + mtp_loss_scaling_factor=0.3, # 0.3 for first 15T tokens, 0.1 after + # Recompute for memory efficiency + recompute_granularity="selective", +) +``` + +### Finetuning Examples + +#### Full Finetuning (GLM 4.5 Air) +```python +from megatron.bridge.recipes.glm import glm45_air_106b_finetune_config + +config = glm45_air_106b_finetune_config( + name="glm45_air_full_finetune", + pretrained_checkpoint="/models/glm45-air-106b", + peft=None, # Full finetuning + train_iters=1000, + global_batch_size=128, + micro_batch_size=1, + finetune_lr=5e-6, +) +``` + +#### LoRA Finetuning (GLM 4.5 Air) +```python +from megatron.bridge.recipes.glm import glm45_air_106b_finetune_config + +config = glm45_air_106b_finetune_config( + name="glm45_air_lora_finetune", + pretrained_checkpoint="/models/glm45-air-106b", + peft="lora", # or "dora" + train_iters=1000, + global_batch_size=128, + micro_batch_size=1, + finetune_lr=1e-4, + # Uses TP=1, PP=2, EP=4 (8 GPUs) automatically +) +``` + +#### DoRA Finetuning (GLM 4.5 355B) +```python +from megatron.bridge.recipes.glm import glm45_355b_finetune_config + +config = glm45_355b_finetune_config( + name="glm45_355b_dora_finetune", + pretrained_checkpoint="/models/glm45-355b", + peft="dora", + train_iters=1000, + global_batch_size=128, + micro_batch_size=1, + finetune_lr=1e-4, + # Uses TP=2, PP=4, EP=4 (32 GPUs) automatically +) +``` + +### Command-Line Training + +```bash +# GLM 4.5 Air - LoRA finetuning on single node (8 GPUs) +torchrun --nproc-per-node=8 run/run_recipe.py \ +--pretrained-checkpoint /models/glm45-air-106b \ +--recipe glm45_air_106b_finetune_config \ +peft=lora \ +train.global_batch_size=128 \ +train.train_iters=1000 \ +checkpoint.save=$SAVE_DIR/glm45_air_lora + +# GLM 4.5 355B - Full finetuning (256 GPUs) +torchrun --nnodes=32 --nproc-per-node=8 run/run_recipe.py \ +--pretrained-checkpoint /models/glm45-355b \ +--recipe glm45_355b_finetune_config \ +peft=None \ +train.global_batch_size=256 \ +train.train_iters=1000 \ +checkpoint.save=$SAVE_DIR/glm45_355b_full +``` + +## Advanced Configuration + +### Multi-Token Prediction (MTP) +MTP can be configured for improved training efficiency: + +```python +config = glm45_355b_pretrain_config( + name="glm45_with_mtp", + mtp_num_layers=1, # Number of MTP prediction layers + mtp_loss_scaling_factor=0.3, # 0.3 early training, 0.1 later + # Set to None or 0 to disable MTP +) +``` + +### Activation Recomputation +For memory-constrained scenarios: + +```python +config = glm45_air_106b_pretrain_config( + name="glm45_with_recompute", + recompute_granularity="selective", # or "full" + recompute_method="uniform", + recompute_num_layers=2, +) +``` + +### Expert Parallelism Tuning +Adjust expert parallelism based on your cluster: + +```python +config = glm45_air_106b_pretrain_config( + name="glm45_custom_parallelism", + tensor_model_parallel_size=1, + pipeline_model_parallel_size=4, + expert_model_parallel_size=16, # Adjust based on GPU count + sequence_parallel=True, +) +``` + +## Examples +- Checkpoint import/export: [examples/conversion/convert_checkpoints.py](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/conversion/convert_checkpoints.py) +- Generate text (HF→Megatron): [examples/conversion/hf_to_megatron_generate_text.py](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/conversion/hf_to_megatron_generate_text.py) + +## Hugging Face Model Cards +- GLM 4.5 355B: https://huggingface.co/zai-org/GLM-4.5 +- GLM 4.5 Air 106B: https://huggingface.co/zai-org/GLM-4.5-Air + +## Key Features for Production + +1. **Efficient Scaling**: MoE architecture enables 355B parameters with only 32B active per token +2. **Single-Node PEFT**: GLM 4.5 Air can be fine-tuned with LoRA/DoRA on just 8 GPUs +3. **Long Context**: Native 131K token context window with optimized RoPE +4. **Multi-Token Prediction**: Faster convergence through MTP training +5. **Flexible Deployment**: Multiple parallelism strategies for different hardware configurations +6. **Load Balancing**: Shared expert overlap and auxiliary loss for optimal expert utilization + +## Related Docs +- Recipe usage: [Recipe usage](/../../recipe-usage) +- Customizing the training recipe configuration: [Configuration overview](/../../training/config-container-overview) +- Training entry points: [Entry points](/../../training/entry-points) diff --git a/fern/v0.2.0/pages/models/llm/gpt-oss.mdx b/fern/v0.2.0/pages/models/llm/gpt-oss.mdx new file mode 100644 index 0000000000..4a9f6e88c4 --- /dev/null +++ b/fern/v0.2.0/pages/models/llm/gpt-oss.mdx @@ -0,0 +1,239 @@ +--- +title: "GPT OSS" +description: "" +--- + +GPT OSS is a Mixture-of-Experts (MoE) language model family featuring two variants: **GPT OSS 20B** and **GPT OSS 120B**. These models are designed with advanced attention mechanisms and MoE architectures optimized for long-context understanding. + +The GPT OSS models feature decoder-only architectures with routed expert layers, supporting context lengths up to 128K tokens through YaRN position embeddings. Both variants use grouped-query attention and specialized attention mechanisms including sliding window attention with learnable softmax. + +GPT OSS models are supported via the Bridge system with specialized configurations for MoE optimizations and long-context training. + +## Model Architecture + +### GPT OSS 20B +- **Parameters**: 20B total +- **Layers**: 24 decoder layers +- **Experts**: 32 routed experts per layer with top-4 routing +- **Hidden size**: 2880 +- **FFN hidden size**: 2880 (dense layers), 2880 (expert layers) +- **Attention heads**: 64 query heads, 8 key-value groups (GQA) +- **KV channels**: 64 +- **Vocab size**: 201,088 +- **Context Length**: 128K tokens (via YaRN) +- **Activation**: QuickGELU with gated linear units +- **Normalization**: RMSNorm + +### GPT OSS 120B +- **Parameters**: 120B total +- **Layers**: 36 decoder layers +- **Experts**: 128 routed experts per layer with top-4 routing +- **Hidden size**: 2880 +- **FFN hidden size**: 2880 (dense layers), 2880 (expert layers) +- **Attention heads**: 64 query heads, 8 key-value groups (GQA) +- **KV channels**: 64 +- **Vocab size**: 201,088 +- **Context Length**: 128K tokens (via YaRN) +- **Activation**: QuickGELU with gated linear units +- **Normalization**: RMSNorm + +## Key Features + +- **YaRN Position Embeddings**: Advanced rotary position embeddings with scaling factor 32.0 for long-context extension +- **Grouped-Query Attention (GQA)**: Efficient attention with 8 key-value groups +- **Sliding Window Attention**: Window size of 128 tokens with alternating full/windowed attention pattern +- **Learnable Softmax**: Novel softmax implementation with learnable offset parameters (sink attention) +- **QuickGELU Activation**: Fast approximate GELU with clamping at 7.0 for stability +- **MoE Routing**: Top-4 expert routing without load balancing loss +- **Grouped GEMM**: Optimized grouped matrix multiplications for expert computation +- **Bias in Linear Layers**: Linear layers include bias terms +- **Activation Clamping**: Output activations clamped to [-7.0, 7.0] for numerical stability + +## Conversion with 🤗 Hugging Face + +### Load HF → Megatron +```python +from megatron.bridge import AutoBridge + +# Example: GPT OSS 20B +bridge = AutoBridge.from_hf_pretrained("openai/gpt-oss-20b") +provider = bridge.to_megatron_provider() + +# Configure parallelism before instantiating the model +provider.tensor_model_parallel_size = 2 +provider.pipeline_model_parallel_size = 4 +provider.expert_model_parallel_size = 4 +provider.sequence_parallel = True + +model = provider.provide_distributed_model(wrap_with_ddp=False) + +provider.finalize() + +# For GPT OSS 120B: +# bridge = AutoBridge.from_hf_pretrained("openai/gpt-oss-120b") +``` + +### Export Megatron → HF +```python +# Convert from a Megatron checkpoint directory to HF format +bridge.export_ckpt( + megatron_path="/results/gpt_oss_120b/checkpoints/iter_0500000", + hf_path="./gpt-oss-hf-export", +) +``` + +## Examples + +- Checkpoint conversion: [examples/conversion/convert_checkpoints.py](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/conversion/convert_checkpoints.py) + +## Pretrain recipes + +### Example usage (GPT OSS 20B) +```python +from megatron.bridge.recipes.gpt_oss import gpt_oss_20b_pretrain_config + +cfg = gpt_oss_20b_pretrain_config( + name="gpt_oss_20b_pretrain", + data_paths=["/path/to/dataset.nvjsonl"], + dir="/results/gpt_oss_20b", + train_iters=500_000, + global_batch_size=512, + seq_length=4096, +) +``` + +### Example usage (GPT OSS 120B) +```python +from megatron.bridge.recipes.gpt_oss import gpt_oss_120b_pretrain_config + +cfg = gpt_oss_120b_pretrain_config( + name="gpt_oss_120b_pretrain", + data_paths=["/path/to/dataset.nvjsonl"], + dir="/results/gpt_oss_120b", + train_iters=500_000, + global_batch_size=512, + seq_length=4096, +) +``` + +### Key configuration options +- **Parallelism (20B)**: Default TP=2, PP=4, EP=4 for efficient MoE training +- **Parallelism (120B)**: Default TP=2, PP=4, EP=16 for large-scale training +- **Sequence parallel**: Enabled by default for memory efficiency +- **Context parallel**: Supports CP for long sequences (131K tokens) +- **Manual GC**: Aggressive garbage collection (interval=100) for stable memory usage +- **MoE optimizations**: Grouped GEMM and permute fusion enabled by default + +## Finetuning recipes + +### Example usage (GPT OSS 20B - LoRA finetuning) +```python +from megatron.bridge.recipes.gpt_oss import gpt_oss_20b_finetune_config + +cfg = gpt_oss_20b_finetune_config( + hf_path="openai/gpt-oss-20b", + name="gpt_oss_20b_lora_finetune", + pretrained_checkpoint="path/to/gpt_oss/checkpoint", + peft="lora", # or "dora" for DoRA + train_iters=1000, + global_batch_size=128, + finetune_lr=1e-4, +) +``` + +### Example usage (GPT OSS 20B - Full SFT) +```python +cfg = gpt_oss_20b_finetune_config( + hf_path="openai/gpt-oss-20b", + name="gpt_oss_20b_full_sft", + pretrained_checkpoint="path/to/gpt_oss/checkpoint", + peft=None, # Full supervised finetuning + train_iters=1000, + global_batch_size=128, + finetune_lr=5e-6, # Lower LR for full SFT +) +``` + +### Example usage (GPT OSS 120B - LoRA finetuning) +```python +from megatron.bridge.recipes.gpt_oss import gpt_oss_120b_finetune_config + +cfg = gpt_oss_120b_finetune_config( + hf_path="openai/gpt-oss-120b", + name="gpt_oss_120b_lora_finetune", + pretrained_checkpoint="path/to/gpt_oss/checkpoint", + peft="lora", + train_iters=1000, + global_batch_size=128, + finetune_lr=1e-4, +) +``` + +### Default configurations + +#### GPT OSS 20B + +**LoRA/DoRA (1 node, 8 GPUs)** +- TP=1, PP=1, EP=1, LR=1e-4 +- Optimized for parameter-efficient training +- Lower memory footprint + +**Full SFT (1 node, 8 GPUs)** +- TP=1, PP=1, EP=8, LR=5e-6 +- Full model training with expert parallelism +- Higher throughput with distributed experts + +#### GPT OSS 120B + +**LoRA/DoRA (2 nodes, 16 GPUs)** +- TP=1, PP=4, EP=8, LR=1e-4 +- Optimized for parameter-efficient training +- Pipeline parallelism for memory efficiency + +**Full SFT (2 nodes, 16 GPUs)** +- TP=1, PP=1, EP=8, LR=5e-6 +- Full model training with expert parallelism +- Requires more memory per GPU + +## API reference + +- GPT OSS recipes: [bridge.recipes.gpt_oss](/../../apidocs/bridge/bridge.recipes.gpt_oss) +- GPT OSS model provider: [bridge.models.gpt_oss.GPTOSSProvider](/../../apidocs/bridge/bridge.models.gpt_oss) + +## Performance optimizations + +### Memory efficiency +- **Manual GC**: Aggressive garbage collection (interval=100) for stable memory usage +- **Precision-aware optimizer**: BF16 gradients and optimizer states +- **Expert parallelism**: Distributes experts across GPUs (EP=4 for 20B, EP=16 for 120B) +- **Sequence parallel**: Reduces activation memory across tensor parallel ranks +- **Context parallel**: Splits long sequences across multiple GPUs + +### Compute efficiency +- **MoE permute fusion**: Fuses expert permutation operations +- **Grouped GEMM**: Optimized expert computation with grouped matrix multiplications +- **AllToAll dispatcher**: Efficient token routing across expert parallel ranks +- **Bias activation fusion**: Fuses bias addition with activation functions +- **Gradient overlapping**: Overlaps gradient all-reduce with backward computation + +### Attention optimizations +- **Sliding window attention**: Reduces attention computation with 128-token windows +- **Window attention skip frequency**: Alternates between windowed and full attention every 2 layers +- **Sink attention**: Learnable softmax offsets for improved attention stability +- **Flash Attention**: FlashAttention-2 support via Transformer Engine +- **Activation clamping**: Prevents numerical instability with [-7.0, 7.0] clamping + +## Hugging Face model cards + +### GPT OSS 20B +- Base: [openai/gpt-oss-20b](https://huggingface.co/openai/gpt-oss-20b) + +### GPT OSS 120B +- Base: [openai/gpt-oss-120b](https://huggingface.co/openai/gpt-oss-120b) + +## Related docs + +- Recipe usage and customization: [Recipe usage](/../../recipe-usage) +- Training configuration: [Configuration overview](/../../training/config-container-overview) +- Training entry points: [Entry points](/../../training/entry-points) +- Attention optimizations: [Attention optimizations](/../../training/attention-optimizations) diff --git a/fern/v0.2.0/pages/models/llm/index.mdx b/fern/v0.2.0/pages/models/llm/index.mdx new file mode 100644 index 0000000000..8a858617d4 --- /dev/null +++ b/fern/v0.2.0/pages/models/llm/index.mdx @@ -0,0 +1,6 @@ +--- +title: "Large Language Models" +description: "" +--- + +This section documents Large Language Models supported by Megatron Bridge, with examples for converting to/from 🤗 Hugging Face and links to training recipes. diff --git a/fern/v0.2.0/pages/models/llm/llama-nemotron.mdx b/fern/v0.2.0/pages/models/llm/llama-nemotron.mdx new file mode 100644 index 0000000000..aa36574f39 --- /dev/null +++ b/fern/v0.2.0/pages/models/llm/llama-nemotron.mdx @@ -0,0 +1,125 @@ +--- +title: "Llama Nemotron" +description: "" +--- + +[Llama Nemotron](https://huggingface.co/collections/nvidia/llama-nemotron) is NVIDIA's family of large language models derived from Meta's Llama architecture, post-trained for enhanced reasoning, human chat preferences, and agentic tasks such as RAG and tool calling. The models feature neural architecture search (NAS) optimizations for improved efficiency and accuracy trade-offs. + +Llama Nemotron models are supported via the Bridge system with auto-detected configuration and weight mapping. + +## Available Models + +Megatron Bridge supports the following Llama Nemotron model variants: + +- **Llama-3.3-Nemotron-Super-49B**: 49B parameters (NAS-optimized from 70B) +- **Llama-3.1-Nemotron-Ultra-253B**: 253B parameters (large-scale reasoning model) +- **Llama-3.1-Nemotron-70B**: 70B parameters (standard size) +- **Llama-3.1-Nemotron-Nano-8B**: 8B parameters (efficient variant) +- **Llama-3.1-Nemotron-Nano-4B**: 4B parameters (ultra-compact variant) + +All models are ready for commercial use and support context lengths up to 128K tokens. + +## Model Architecture Features + +- **Neural Architecture Search (NAS)**: Novel approach to reduce memory footprint while maintaining accuracy +- **Heterogeneous Blocks**: Non-standard and non-repetitive layer configurations for efficiency + - Skip attention in some blocks + - Variable FFN expansion/compression ratios between blocks +- **Multi-Phase Post-Training**: + - Supervised fine-tuning for Math, Code, Science, and Tool Calling + - Reward-aware Preference Optimization (RPO) for chat + - Reinforcement Learning with Verifiable Rewards (RLVR) for reasoning + - Iterative Direct Preference Optimization (DPO) for tool calling +- **Extended Context**: Native support for sequences up to 128K tokens +- **Commercial Ready**: Fully licensed for commercial deployment + +## Conversion with 🤗 Hugging Face + +### Load HF → Megatron + +```python +from megatron.bridge import AutoBridge + +# Example: Llama-3.3-Nemotron-Super-49B +bridge = AutoBridge.from_hf_pretrained( + "nvidia/Llama-3_3-Nemotron-Super-49B-v1_5", + trust_remote_code=True +) +provider = bridge.to_megatron_provider() + +# Optionally configure parallelism before instantiating the model +provider.tensor_model_parallel_size = 2 +provider.pipeline_model_parallel_size = 1 + +model = provider.provide_distributed_model(wrap_with_ddp=False) +``` + +**Note**: Heterogeneous Llama-Nemotron models (Super/Ultra) require `trust_remote_code=True` as they use custom `DeciLMForCausalLM` architecture. Homogeneous models (Nano/70B) use standard Llama architecture and don't require this flag. + +### Import Checkpoint from HF + +```bash +python examples/conversion/convert_checkpoints.py import \ + --hf-model nvidia/Llama-3_3-Nemotron-Super-49B-v1_5 \ + --megatron-path /checkpoints/llama_nemotron_super_49b_megatron \ + --trust-remote-code +``` + +### Export Megatron → HF + +```python +from megatron.bridge import AutoBridge + +# Load the bridge from HF model ID +bridge = AutoBridge.from_hf_pretrained( + "nvidia/Llama-3_3-Nemotron-Super-49B-v1_5", + trust_remote_code=True +) + +# Export a trained Megatron checkpoint to HF format +bridge.export_ckpt( + megatron_path="/results/llama_nemotron_super_49b/checkpoints/iter_0000500", + hf_path="/exports/llama_nemotron_super_49b_hf", +) +``` + +### Run Inference on Converted Checkpoint + +```bash +python examples/conversion/hf_to_megatron_generate_text.py \ + --hf_model_path nvidia/Llama-3_3-Nemotron-Super-49B-v1_5 \ + --megatron_model_path /checkpoints/llama_nemotron_super_49b_megatron \ + --prompt "What is artificial intelligence?" \ + --max_new_tokens 100 \ + --tp 2 \ + --trust-remote-code +``` + +For more details, see [examples/conversion/hf_to_megatron_generate_text.py](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/conversion/hf_to_megatron_generate_text.py) + +## Recipes + +Training recipes for Llama Nemotron models are not currently available. + +## Hugging Face Model Cards & References + +### Hugging Face Model Cards +- Llama Nemotron Collection: https://huggingface.co/collections/nvidia/llama-nemotron +- Llama-3.3-Nemotron-Super-49B-v1.5: https://huggingface.co/nvidia/Llama-3_3-Nemotron-Super-49B-v1_5 +- Llama-3.1-Nemotron-Ultra-253B-v1: https://huggingface.co/nvidia/Llama-3_1-Nemotron-Ultra-253B-v1 +- Llama-3.1-Nemotron-Nano-8B-v1: https://huggingface.co/nvidia/Llama-3.1-Nemotron-Nano-8B-v1 +- Llama-3.1-Nemotron-Nano-4B-v1.1: https://huggingface.co/nvidia/Llama-3.1-Nemotron-Nano-4B-v1.1 + +### Technical Papers +- Llama-Nemotron: Efficient Reasoning Models: [arXiv:2505.00949](https://arxiv.org/abs/2505.00949) +- Puzzle: Distillation-Based NAS for Inference-Optimized LLMs: [arXiv:2411.19146](https://arxiv.org/abs/2411.19146) +- Reward-aware Preference Optimization: [arXiv:2502.00203](https://arxiv.org/abs/2502.00203) + +### Additional Resources +- NVIDIA Build Platform: https://build.nvidia.com/ +- Llama Nemotron Post-Training Dataset: https://huggingface.co/nvidia/Llama-Nemotron-Post-Training-Dataset + +## Related Docs +- Recipe usage: [Recipe usage](/../../recipe-usage) +- Customizing the training recipe configuration: [Configuration overview](/../../training/config-container-overview) +- Training entry points: [Entry points](/../../training/entry-points) diff --git a/fern/v0.2.0/pages/models/llm/llama3.mdx b/fern/v0.2.0/pages/models/llm/llama3.mdx new file mode 100644 index 0000000000..428720bab9 --- /dev/null +++ b/fern/v0.2.0/pages/models/llm/llama3.mdx @@ -0,0 +1,228 @@ +--- +title: "Llama 3" +description: "" +--- + +[Meta's Llama](https://www.llama.com/models/llama-3/) builds on the general transformer decoder framework with some key additions such as pre-normalization, SwiGLU activations, and Rotary Positional Embeddings (RoPE). More information is available in the companion paper ["Llama: Open and Efficient Foundation Language Models"](https://arxiv.org/abs/2302.13971). With a wide variety of model sizes - Llama has options for every inference budget. + +Llama family models are supported via the Bridge system with auto-detected configuration and weight mapping. + +## Available Models + +Megatron Bridge supports the following Llama model variants: + +- **Llama 3.2**: 1B, 3B +- **Llama 3**: 8B, 70B (with 8K, 16K, 64K, 128K context variants) +- **Llama 3.1**: 8B, 70B, 405B (with 128K context length) + +All models support both pretraining and finetuning with full parameter updates or PEFT methods (LoRA, DoRA). + +## Model Architecture Features + +- **Pre-normalization**: RMSNorm before each transformer sub-layer for training stability +- **SwiGLU Activation**: Gated linear units in the feedforward network +- **Rotary Positional Embeddings (RoPE)**: Relative position encoding via rotation matrices +- **Grouped Query Attention (GQA)**: Memory-efficient attention mechanism (70B+ models) +- **Extended Context**: Native support for long sequences up to 128K tokens (Llama 3.1) + +## Conversion with 🤗 Hugging Face + +### Load HF → Megatron + +```python +from megatron.bridge import AutoBridge + +# Example: Llama 3.1 8B +bridge = AutoBridge.from_hf_pretrained("meta-llama/Meta-Llama-3.1-8B") +provider = bridge.to_megatron_provider() + +# Optionally configure parallelism before instantiating the model +provider.tensor_model_parallel_size = 2 +provider.pipeline_model_parallel_size = 1 + +model = provider.provide_distributed_model(wrap_with_ddp=False) +``` + +### Import Checkpoint from HF + +```bash +python examples/conversion/convert_checkpoints.py import \ + --hf-model meta-llama/Meta-Llama-3.1-8B \ + --megatron-path /checkpoints/llama31_8b_megatron +``` + +### Export Megatron → HF + +```python +from megatron.bridge import AutoBridge + +# Load the bridge from HF model ID +bridge = AutoBridge.from_hf_pretrained("meta-llama/Meta-Llama-3.1-8B") + +# Export a trained/finetuned Megatron checkpoint to HF format +bridge.export_ckpt( + megatron_path="/results/llama31_8b/checkpoints/iter_0000500", + hf_path="/exports/llama31_8b_hf", +) +``` + +### Run Inference on Converted Checkpoint + +```bash +python examples/conversion/hf_to_megatron_generate_text.py \ + --hf_model_path meta-llama/Meta-Llama-3.1-8B \ + --megatron_model_path /checkpoints/llama31_8b_megatron \ + --prompt "What is artificial intelligence?" \ + --max_new_tokens 100 \ + --tp 2 +``` + +For more details, see [examples/conversion/hf_to_megatron_generate_text.py](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/conversion/hf_to_megatron_generate_text.py) + +## Recipes + +See: [bridge.recipes.llama.llama3](/../../apidocs/bridge/bridge.recipes.llama.llama3) + +### Available Recipes + +- **Pretrain recipes**: + - `llama32_1b_pretrain_config`, `llama32_3b_pretrain_config`: Llama 3.2 (1B, 3B) + - `llama3_8b_pretrain_config`: Llama 3 8B with 8K context + - `llama3_8b_16k_pretrain_config`, `llama3_8b_64k_pretrain_config`, `llama3_8b_128k_pretrain_config`: Llama 3 8B with extended context (16K/64K/128K) + - `llama3_8b_low_precision_pretrain_config`: Llama 3 8B with low precision (FP8/MXFP8/NVFP4) + - `llama3_70b_pretrain_config`, `llama3_70b_16k_pretrain_config`, `llama3_70b_64k_pretrain_config`: Llama 3 70B (8K/16K/64K context) + - `llama31_8b_pretrain_config`, `llama31_70b_pretrain_config`, `llama31_405b_pretrain_config`: Llama 3.1 (8B/70B/405B, 128K context) + +- **Finetune recipes**: + - `llama32_1b_finetune_config`, `llama32_3b_finetune_config`: Llama 3.2 with PEFT support + - `llama3_8b_finetune_config`, `llama31_8b_finetune_config`: Llama 3/3.1 8B with PEFT support + - `llama3_70b_finetune_config`, `llama31_70b_finetune_config`: Llama 3/3.1 70B with PEFT support + - `llama31_405b_finetune_config`: Llama 3.1 405B with PEFT support + +### Parallelism Configurations + +#### Llama 3.2 (1B, 3B) +| Model | Mode | TP | PP | Total GPUs | Use Case | +|-------|------|----|----|------------|----------| +| **1B / 3B** | Pretrain | 1 | 1 | 8 | Pre-training (single node) | +| **1B / 3B** | Full SFT | 1 | 1 | 8 | Full supervised finetuning | +| **1B / 3B** | LoRA/DoRA | 1 | 1 | 8 | PEFT finetuning | + +#### Llama 3 / 3.1 (8B) +| Model | Mode | TP | PP | CP | Total GPUs | Use Case | +|-------|------|----|----|----|-----------:|----------| +| **8B** | Pretrain | 1 | 1 | 2 | 16 | Pre-training | +| **8B** | Full SFT | 2 | 1 | 1 | 16 | Full supervised finetuning | +| **8B** | LoRA/DoRA | 1 | 1 | 1 | 8 | PEFT finetuning (single node) | + +#### Llama 3 / 3.1 (70B) +| Model | Mode | TP | PP | VP | CP | Total GPUs | Use Case | +|-------|------|----|----|----|----|------------|----------| +| **70B** | Pretrain | 4 | 4 | 5 | 2 | 64 | Pre-training | +| **70B** | Full SFT | 8 | 4 | - | 1 | 256 | Full supervised finetuning (32 nodes) | +| **70B** | LoRA/DoRA | 8 | 1 | - | 1 | 8 | PEFT finetuning (single node!) | + +#### Llama 3.1 (405B) +| Model | Mode | TP | PP | VP | CP | Total GPUs | Use Case | +|-------|------|----|----|----|----|------------|----------| +| **405B** | Pretrain | 8 | 8 | 2 | 4 | 512 | Pre-training (64 nodes) | +| **405B** | Full SFT | 8 | 16 | - | 1 | 2048 | Full supervised finetuning (256 nodes) | +| **405B** | LoRA/DoRA | 4 | 8 | 8 | 1 | 256 | PEFT finetuning (32 nodes) | + +**Key Features**: +- **Context Parallelism**: Enabled for long context training (16K/64K/128K variants) +- **Sequence Parallel**: Enabled by default for larger models (70B+) for memory efficiency +- **Low Precision Training**: FP8, MXFP8, NVFP4 options available for 8B model +- **Virtual Pipeline**: VP parallelism for 70B and 405B models + +### Pre-training Example + +```python +from megatron.bridge.recipes.llama import llama3_8b_pretrain_config + +config = llama3_8b_pretrain_config( + name="llama3_8b_pretrain", + data_paths=["/path/to/dataset.nvjsonl"], + dir="/results/llama3_8b", + train_iters=500_000, + global_batch_size=512, + seq_length=8192, + # Uses TP=1, PP=1, CP=2 (16 GPUs) automatically +) +``` + +### Finetuning Examples + +**Before finetuning**, ensure these environment variables are set: +- `SAVE_DIR`: checkpoint and log saving directory +- `HF_TOKEN`: to download models from HF Hub (if required) +- `HF_HOME`: (optional) to avoid re-downloading models and datasets +- `WANDB_API_KEY`: (optional) to enable WandB logging + +#### Full Finetuning (Llama 3 8B) + +```python +from megatron.bridge.recipes.llama import llama3_8b_finetune_config + +cfg = llama3_8b_finetune_config( + name="llama3_8b_full_sft", + pretrained_checkpoint="/results/llama3_8b/checkpoints/iter_0500000", + peft=None, # Full supervised finetuning + train_iters=1000, + global_batch_size=64, + finetune_lr=5e-6, + # Uses TP=2, PP=1 (16 GPUs) automatically +) +``` + +#### LoRA Finetuning + +```python +from megatron.bridge.recipes.llama import llama3_8b_finetune_config + +cfg = llama3_8b_finetune_config( + name="llama3_8b_lora", + pretrained_checkpoint="/results/llama3_8b/checkpoints/iter_0500000", + peft="lora", # or "dora" for DoRA + train_iters=1000, + global_batch_size=128, + finetune_lr=1e-4, + # Uses TP=1, PP=1 (8 GPUs) automatically +) +``` + +#### LoRA Finetuning + +```python +from megatron.bridge.recipes.llama import llama3_70b_finetune_config + +cfg = llama3_70b_finetune_config( + name="llama3_70b_lora", + pretrained_checkpoint="/results/llama3_70b/checkpoints/iter_0500000", + peft="lora", + train_iters=1000, + global_batch_size=128, + finetune_lr=1e-4, + # Uses TP=8, PP=1 (8 GPUs) automatically +) +``` + +## Hugging Face Model Cards & References + +### Hugging Face Model Cards +- Llama 3.2 1B: https://huggingface.co/meta-llama/Llama-3.2-1B +- Llama 3.2 3B: https://huggingface.co/meta-llama/Llama-3.2-3B +- Llama 3 8B: https://huggingface.co/meta-llama/Meta-Llama-3-8B +- Llama 3 70B: https://huggingface.co/meta-llama/Meta-Llama-3-70B +- Llama 3.1 8B: https://huggingface.co/meta-llama/Meta-Llama-3.1-8B +- Llama 3.1 70B: https://huggingface.co/meta-llama/Meta-Llama-3.1-70B +- Llama 3.1 405B: https://huggingface.co/meta-llama/Meta-Llama-3.1-405B + +### Technical Papers +- Llama: Open and Efficient Foundation Language Models: [arXiv:2302.13971](https://arxiv.org/abs/2302.13971) +- The Llama 3 Herd of Models: [arXiv:2407.21783](https://arxiv.org/abs/2407.21783) + +## Related Docs +- Recipe usage: [Recipe usage](/../../recipe-usage) +- Customizing the training recipe configuration: [Configuration overview](/../../training/config-container-overview) +- Training entry points: [Entry points](/../../training/entry-points) diff --git a/fern/v0.2.0/pages/models/llm/mistral.mdx b/fern/v0.2.0/pages/models/llm/mistral.mdx new file mode 100644 index 0000000000..06bb18635a --- /dev/null +++ b/fern/v0.2.0/pages/models/llm/mistral.mdx @@ -0,0 +1,105 @@ +--- +title: "Mistral" +description: "" +--- + +[Mistral AI](https://mistral.ai/) develops frontier large language models with a focus on efficiency and performance. The Mistral family includes both dense and Mixture-of-Experts architectures, featuring innovations like sliding window attention and efficient context handling. + +Mistral models are supported via the Bridge system with auto-detected configuration and weight mapping. + +## Available Models + +Megatron Bridge supports the following Mistral model variants: + +- **Mistral Small 3 (24B)**: 24B parameters with 128K context length +- **Mistral 7B**: 7B parameters, efficient baseline model +- **Mistral 7B Instruct**: Instruction-tuned variant + +Additional Mistral models (including MoE variants like Mixtral) may be supported through the standard conversion pipeline. + +## Model Architecture Features + +- **Sliding Window Attention**: Efficient attention mechanism for long sequences +- **Grouped Query Attention (GQA)**: Memory-efficient attention mechanism +- **Rotary Positional Embeddings (RoPE)**: Relative position encoding +- **SwiGLU Activation**: Gated linear units in the feedforward network +- **Extended Context**: Support for sequences up to 128K tokens (Mistral Small 3) +- **YaRN RoPE Scaling**: Advanced rope scaling for extended context lengths + +## Conversion with 🤗 Hugging Face + +### Load HF → Megatron + +```python +from megatron.bridge import AutoBridge + +# Example: Mistral Small 3 24B +bridge = AutoBridge.from_hf_pretrained("mistralai/Mistral-Small-24B-Base-2501") +provider = bridge.to_megatron_provider() + +# Optionally configure parallelism before instantiating the model +provider.tensor_model_parallel_size = 2 +provider.pipeline_model_parallel_size = 1 + +model = provider.provide_distributed_model(wrap_with_ddp=False) +``` + +### Import Checkpoint from HF + +```bash +python examples/conversion/convert_checkpoints.py import \ + --hf-model mistralai/Mistral-Small-24B-Base-2501 \ + --megatron-path /checkpoints/mistral_small_24b_megatron +``` + +### Export Megatron → HF + +```python +from megatron.bridge import AutoBridge + +# Load the bridge from HF model ID +bridge = AutoBridge.from_hf_pretrained("mistralai/Mistral-Small-24B-Base-2501") + +# Export a trained Megatron checkpoint to HF format +bridge.export_ckpt( + megatron_path="/results/mistral_small_24b/checkpoints/iter_0000500", + hf_path="/exports/mistral_small_24b_hf", +) +``` + +### Run Inference on Converted Checkpoint + +```bash +python examples/conversion/hf_to_megatron_generate_text.py \ + --hf_model_path mistralai/Mistral-Small-24B-Base-2501 \ + --megatron_model_path /checkpoints/mistral_small_24b_megatron \ + --prompt "What is artificial intelligence?" \ + --max_new_tokens 100 \ + --tp 2 +``` + +For more details, see [examples/conversion/hf_to_megatron_generate_text.py](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/conversion/hf_to_megatron_generate_text.py) + +## Recipes + +Training recipes for Mistral models are not currently available. The Bridge supports checkpoint conversion for inference and deployment use cases. + +## Hugging Face Model Cards & References + +### Hugging Face Model Cards +- Mistral Small 3 (24B): https://huggingface.co/mistralai/Mistral-Small-24B-Base-2501 +- Mistral Small 3 (24B) Instruct: https://huggingface.co/mistralai/Mistral-Small-24B-Instruct-2501 +- Mistral 7B v0.1: https://huggingface.co/mistralai/Mistral-7B-v0.1 +- Mistral 7B Instruct v0.2: https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.2 + +### Technical Papers +- Mistral 7B: https://arxiv.org/abs/2310.06825 + +### Additional Resources +- Mistral AI Website: https://mistral.ai/ +- Mistral Documentation: https://docs.mistral.ai/ + +## Related Docs +- Recipe usage: [Recipe usage](/../../recipe-usage) +- Customizing the training recipe configuration: [Configuration overview](/../../training/config-container-overview) +- Training entry points: [Entry points](/../../training/entry-points) diff --git a/fern/v0.2.0/pages/models/llm/moonlight.mdx b/fern/v0.2.0/pages/models/llm/moonlight.mdx new file mode 100644 index 0000000000..cac7f2be36 --- /dev/null +++ b/fern/v0.2.0/pages/models/llm/moonlight.mdx @@ -0,0 +1,164 @@ +--- +title: "Moonlight" +description: "" +--- + +[Moonlight](https://huggingface.co/moonshotai/Moonlight-16B-A3B) is a 16B-parameter Mixture-of-Experts (MoE) model from **Moonshot AI** trained with 5.7T tokens using the innovative **Muon optimizer**. While Moonlight shares the same architecture as DeepSeek-V3 (featuring Multi-head Latent Attention and MoE), it is a distinct model that advances the Pareto frontier of performance vs training FLOPs through the use of Muon, which is ~2× more sample efficient than Adam with compute optimal training. + +The model features 27 decoder layers with 64 routed experts and 8 shared experts per layer, with 3B activated parameters per forward pass out of 16B total parameters. + +Moonlight models are supported via the Bridge system with specialized configurations for MoE and MLA optimizations. + +## Model Architecture + +- **Parameters**: 16B total, 3B activated per forward pass +- **Layers**: 27 decoder layers +- **Attention**: Multi-head Latent Attention (MLA) with RoPE fusion support +- **MoE**: 64 routed experts + 8 shared experts per layer +- **Hidden size**: 2048 +- **Intermediate size**: 10944 (with MLP and expert gating) +- **Vocab size**: 151,936 +- **Context Length**: 8K tokens +- **Training**: 5.7T tokens with Muon optimizer + +## Conversion with 🤗 Hugging Face + +Moonlight shares the same architecture as DeepSeek-V3, which enables compatibility with various inference engines like vLLM and SGLang. The model can be loaded from HuggingFace or used with Megatron checkpoints. + +### Load HF → Megatron +```python +from megatron.bridge import AutoBridge + +# Example: Moonlight-16B-A3B +bridge = AutoBridge.from_hf_pretrained("moonshotai/Moonlight-16B-A3B") +provider = bridge.to_megatron_provider() + +# Configure parallelism before instantiating the model +provider.tensor_model_parallel_size = 2 +provider.pipeline_model_parallel_size = 1 +provider.expert_model_parallel_size = 8 +provider.sequence_parallel = True + +provider.finalize() +model = provider.provide_distributed_model(wrap_with_ddp=False) +``` + +### Export Megatron → HF +```python +# Convert from a Megatron checkpoint directory to HF format +bridge.export_ckpt( + megatron_path="/results/moonlight_16b/checkpoints/iter_0500000", + hf_path="./moonlight-hf-export", +) +``` + +## Examples + +- Checkpoint conversion: [examples/conversion/convert_checkpoints.py](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/conversion/convert_checkpoints.py) + +## Recipes + +See: [bridge.recipes.moonlight](/../../apidocs/bridge/bridge.recipes.moonlight) + +### Available Recipes + +- **Pretrain recipes**: + - `moonlight_16b_pretrain_config`: Pre-training for Moonlight-16B (16B parameters, 3B activated per token) + +- **Finetune recipes**: + - `moonlight_16b_finetune_config`: Finetuning for Moonlight-16B with PEFT support (LoRA, DoRA) + +### Parallelism Configurations + +| Model | Mode | TP | PP | EP | Total GPUs | Use Case | +|-------|------|----|----|----|-----------:|----------| +| **Moonlight-16B** | Pretrain | 2 | 1 | 8 | 16 | Pre-training (2 nodes) | +| **Moonlight-16B** | Full SFT | 2 | 1 | 8 | 16 | Full supervised finetuning (2 nodes) | +| **Moonlight-16B** | LoRA/DoRA | 1 | 1 | 1 | 8 | PEFT finetuning (single node!) | + +**Key Features**: +- **Expert Parallelism**: EP=8 for efficient MoE training (64 experts) +- **Sequence Parallel**: Enabled by default for memory efficiency +- **Selective Recomputation**: Reduces activation memory +- **RoPE Fusion**: Optional MLA-specific optimization (`apply_rope_fusion=True`) +- **DeePEP**: Optional expert permutation optimization (`enable_deepep=True`) + +**Performance Optimizations**: +- **MoE Permute Fusion**: Fused expert permutation operations +- **RoPE Fusion**: Optional fusion for Multi-head Latent Attention +- **Manual GC**: Aggressive garbage collection (interval=5) +- **Precision-Aware Optimizer**: BF16 gradients and optimizer states with FP32 master weights + +**Pipeline Layouts** (optional): +- **PP=1**: No pipelining (default) +- **PP=2**: 14+13 layer split with embedding/loss +- **PP=4**: 8+7+7+6 layer split +- **PP=8**: 5+4+4+4+4+4+4+4 layer split +- **VP**: PP=2,VP=2 and PP=4,VP=2 supported + +### Pre-training Example + +```python +from megatron.bridge.recipes.moonlight import moonlight_16b_pretrain_config + +cfg = moonlight_16b_pretrain_config( + name="moonlight_pretrain", + data_paths=["/path/to/dataset.nvjsonl"], + dir="/results/moonlight_16b", + train_iters=500_000, + global_batch_size=2048, + seq_length=4096, + # Uses TP=2, PP=1, EP=8 (16 GPUs) automatically +) +``` + +### Finetuning Examples + +#### Full Finetuning (2 Nodes) + +```python +from megatron.bridge.recipes.moonlight import moonlight_16b_finetune_config + +cfg = moonlight_16b_finetune_config( + tokenizer_path="moonshotai/Moonlight-16B-A3B", + name="moonlight_full_sft", + pretrained_checkpoint="/results/moonlight_16b/checkpoints/iter_0500000", + peft=None, # Full supervised finetuning + train_iters=1000, + global_batch_size=128, + finetune_lr=5e-6, + # Uses TP=2, PP=1, EP=8 (16 GPUs) automatically +) +``` + +#### LoRA Finetuning + +```python +from megatron.bridge.recipes.moonlight import moonlight_16b_finetune_config + +cfg = moonlight_16b_finetune_config( + tokenizer_path="moonshotai/Moonlight-16B-A3B", + name="moonlight_lora_finetune", + pretrained_checkpoint="/results/moonlight_16b/checkpoints/iter_0500000", + peft="lora", # or "dora" for DoRA + train_iters=1000, + global_batch_size=128, + finetune_lr=1e-4, + # Uses TP=1, PP=1, EP=1 (8 GPUs) automatically +) +``` + +## Hugging Face model cards + +- Moonlight-16B-A3B (Base): [moonshotai/Moonlight-16B-A3B](https://huggingface.co/moonshotai/Moonlight-16B-A3B) +- Moonlight-16B-A3B-Instruct: [moonshotai/Moonlight-16B-A3B-Instruct](https://huggingface.co/moonshotai/Moonlight-16B-A3B-Instruct) + +## Technical Paper + +- Muon is Scalable for LLM Training: [arXiv:2502.16982](https://arxiv.org/abs/2502.16982) + +## Related docs + +- Recipe usage and customization: [Recipe usage](/../../recipe-usage) +- Training configuration: [Configuration overview](/../../training/config-container-overview) +- Training entry points: [Entry points](/../../training/entry-points) diff --git a/fern/v0.2.0/pages/models/llm/nemotron3.mdx b/fern/v0.2.0/pages/models/llm/nemotron3.mdx new file mode 100644 index 0000000000..d681f136ee --- /dev/null +++ b/fern/v0.2.0/pages/models/llm/nemotron3.mdx @@ -0,0 +1,97 @@ +--- +title: "Nemotron 3 Nano" +description: "" +--- + +[Nemotron 3 Nano](https://huggingface.co/collections/nvidia/nvidia-nemotron-v3) is a large language model (LLM) trained from scratch by NVIDIA, and designed as a unified model for both reasoning and non-reasoning tasks. The model employs a hybrid Mixture-of-Experts (MoE) architecture, consisting of 23 Mamba-2 and MoE layers, along with 6 Attention layers. Each MoE layer includes 128 experts plus 1 shared expert, with 5 experts activated per token. The model has 3.5B active parameters and 30B parameters in total. + +NeMo Megatron Bridge supports pretraining, full parameters finetuning, and LoRA finetuning this model. The finetuned model can be converted back to the 🤗 Hugging Face format for downstream evaluation. + + +Please use the custom container `nvcr.io/nvidia/nemo:25.11.nemotron_3_nano` when working with this model. + +Run all commands from `/opt/Megatron-Bridge` (e.g. `docker run -w /opt/Megatron-Bridge ...`) + + + +We use the following environment variables throughout this page +- `HF_MODEL_ID=nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16` +- `MEGATRON_MODEL_PATH=/models/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16` (feel free to set your own path) + + +## Conversion with 🤗 Hugging Face + +### Import HF → Megatron +To import the HF model to your desired `$MEGATRON_MODEL_PATH`, run the following command. +```bash +python examples/conversion/convert_checkpoints.py import \ +--hf-model $HF_MODEL_ID \ +--megatron-path /path/to/output/megatron/ckpt \ +--trust-remote-code +``` + +### Export Megatron → HF +```bash +python examples/conversion/convert_checkpoints.py export \ +--hf-model $HF_MODEL_ID \ +--megatron-path /path/to/trained/megatron/ckpt \ +--hf-path /path/to/output/hf/ckpt +``` + +## Pretraining Examples +```bash +BLEND_PATH=/path/to/dataset/blend +TOKENIZER_MODEL=/path/to/tiktok/tokenizer/model + +torchrun --nproc-per-node=8 examples/models/nemotron_3/pretrain_nemotron_3_nano.py \ +--per-split-data-args-path=${BLEND_PATH} \ +--tokenizer-model=${TOKENIZER_MODEL} \ +train.global_batch_size=3072 \ +train.train_iters=39500 \ +scheduler.lr_warmup_iters=350 +``` + +Notes: +- The default parallelism settings are TP=4, EP=8, PP=1, CP=1. It is recommended to run this pretraining on 4 H100 nodes (32 GPUs). +- To enable wandb logging, you can append `logger.wandb_project=PROJECT_NAME`, `wandb_entity=ENTITY_NAME`, and `wandb_exp_name=EXP_NAME` arguments +- If `BLEND_PATH` and `TOKENIZER_MODEL` are not specified, mock dataset will be used. + +## Finetuning Recipes + +### Full Parameter Fine-Tuning +```bash +torchrun --nproc-per-node=8 examples/models/nemotron_3/finetune_nemotron_3_nano.py \ +train.global_batch_size=128 \ +train.train_iters=100 \ +scheduler.lr_warmup_iters=10 \ +checkpoint.pretrained_checkpoint=/path/to/output/megatron/ckpt +``` + +Notes: +- Default parallelism TP=1, EP=8, PP=1, CP=1. It is recommended to run this recipe on at least 2 H100 nodes (16 GPUs). +- By default, the [SQuAD](https://huggingface.co/datasets/rajpurkar/squad) dataset is used. To use customerized dataset, see this [tutorial](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/tutorials/recipes/llama#quickstart) +- Fine-tuning requires a pretrained megatron checkpoint, which can be obtained in "Import HF → Megatron" section above + +### LoRA Fine-Tuning +To enable LoRA fine-tuning, pass `--peft lora` to script +```bash +torchrun --nproc-per-node=8 examples/models/nemotron_3/finetune_nemotron_3_nano.py \ +--peft lora \ +train.global_batch_size=128 \ +train.train_iters=100 \ +scheduler.lr_warmup_iters=10 \ +checkpoint.pretrained_checkpoint=/path/to/output/megatron/ckpt +``` + +Notes: +- By default, the target modules are linear layers `["linear_qkv", "linear_proj", "linear_fc1", "linear_fc2", "in_proj", "out_proj"]` in the model +- The rest of settings are the same as full parameter fine-tuning above. + +A LoRA checkpoint only contains the learnable adapter weights. In order to convert the LoRA checkpoint to Hugging Face format for downstream evaluation, it is necessary to merge the LoRA adapters back to the base model. + +```bash +python examples/peft/merge_lora.py \ +--hf-model-path $HF_MODEL_ID \ +--lora-checkpoint /path/to/lora/ckpt/iter_xxxxxxx +--output /path/to/merged/ckpt +``` diff --git a/fern/v0.2.0/pages/models/llm/nemotronh.mdx b/fern/v0.2.0/pages/models/llm/nemotronh.mdx new file mode 100644 index 0000000000..5c79ef09e7 --- /dev/null +++ b/fern/v0.2.0/pages/models/llm/nemotronh.mdx @@ -0,0 +1,445 @@ +--- +title: "Nemotron H and Nemotron Nano v2" +description: "" +--- + +[Nemotron H](https://huggingface.co/collections/nvidia/nemotron-h) and [Nemotron Nano v2](https://huggingface.co/collections/nvidia/nvidia-nemotron-v2) are families of **hybrid SSM-Attention models** from **NVIDIA** that combine Mamba (State Space Model) layers with traditional attention layers. These models achieve strong performance while maintaining computational efficiency through their hybrid architecture. + +The Nemotron H family includes models from 4B to 56B parameters with 8K context length, while Nemotron Nano v2 models (9B and 12B) are optimized for edge deployment with extended 128K context support. + +## Model Families + +### Nemotron H +- **4B**: 52 layers, 3072 hidden size, 8K context +- **8B**: 52 layers, 4096 hidden size, 8K context +- **47B**: 98 layers, 8192 hidden size, 8K context +- **56B**: 118 layers, 8192 hidden size, 8K context + +### Nemotron Nano v2 +- **9B**: 56 layers, 4480 hidden size, 128K context +- **12B**: 62 layers, 5120 hidden size, 128K context + +All models are supported via the Bridge system with specialized configurations for hybrid SSM-Attention architecture. + +## Model Architecture + +### Common Features Across All Models +- **Architecture**: Hybrid SSM-Attention (Mamba + Multi-Query Attention) +- **SSM**: Mamba-2 selective state space layers +- **Attention**: Multi-query attention with QK LayerNorm and RoPE +- **Activation**: Squared ReLU (SwiGLU in FFN) +- **Normalization**: RMSNorm +- **Position Embedding**: RoPE (Rotary Position Embeddings) +- **Hybrid Pattern**: Configurable layer-wise mixing of Mamba ("M") and Attention ("*") layers + +### Nemotron H 4B Specifications +- **Parameters**: 4B +- **Layers**: 52 (Hybrid pattern: `M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M-`) +- **Hidden size**: 3072 +- **FFN hidden size**: 12288 +- **Attention heads**: 32 query heads, 8 key-value groups +- **KV channels**: 128 +- **Mamba heads**: 112 +- **Mamba head dim**: 64 +- **Mamba state dim**: 128 +- **Context Length**: 8K tokens + +### Nemotron H 8B Specifications +- **Parameters**: 8B +- **Layers**: 52 (Hybrid pattern: `M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M-`) +- **Hidden size**: 4096 +- **FFN hidden size**: 21504 +- **Attention heads**: 32 query heads, 8 key-value groups +- **KV channels**: 128 +- **Mamba heads**: 128 +- **Mamba head dim**: 64 +- **Mamba state dim**: 128 +- **Context Length**: 8K tokens + +### Nemotron H 47B Specifications +- **Parameters**: 47B +- **Layers**: 98 +- **Hidden size**: 8192 +- **FFN hidden size**: 30720 +- **Attention heads**: 64 query heads, 8 key-value groups +- **KV channels**: 128 +- **Mamba heads**: 256 +- **Mamba head dim**: 64 +- **Mamba state dim**: 256 +- **Context Length**: 8K tokens + +### Nemotron H 56B Specifications +- **Parameters**: 56B +- **Layers**: 118 +- **Hidden size**: 8192 +- **FFN hidden size**: 32768 +- **Attention heads**: 64 query heads, 8 key-value groups +- **KV channels**: 128 +- **Mamba heads**: 256 +- **Mamba head dim**: 64 +- **Mamba state dim**: 256 +- **Context Length**: 8K tokens + +### Nemotron Nano 9B v2 Specifications +- **Parameters**: 9B +- **Layers**: 56 (Hybrid pattern: `M-M-M-MM-M-M-M*-M-M-M*-M-M-M-M*-M-M-M-M*-M-MM-M-M-M-M-M-`) +- **Hidden size**: 4480 +- **FFN hidden size**: 15680 +- **Attention heads**: 40 query heads, 8 key-value groups +- **KV channels**: 128 +- **Mamba heads**: 128 +- **Mamba head dim**: 80 +- **Mamba state dim**: 128 +- **Context Length**: 128K tokens +- **Vocab size**: 131,072 + +### Nemotron Nano 12B v2 Specifications +- **Parameters**: 12B +- **Layers**: 62 (Hybrid pattern: `M-M-M-M*-M-M-M-M*-M-M-M-M*-M-M-M-M*-M-M-M-M*-M-M-M-M*-M-M-M-M-`) +- **Hidden size**: 5120 +- **FFN hidden size**: 20480 +- **Attention heads**: 40 query heads, 8 key-value groups +- **KV channels**: 128 +- **Mamba heads**: 128 +- **Mamba head dim**: 80 +- **Mamba state dim**: 128 +- **Context Length**: 128K tokens +- **Vocab size**: 131,072 + +## Key Features + +### Hybrid SSM-Attention Architecture +- **Mamba Layers (M)**: State space model layers for efficient long-range modeling +- **Attention Layers (*)**: Standard multi-query attention for complex reasoning +- **Configurable Pattern**: Each model has a predefined hybrid pattern balancing efficiency and performance + +### Advanced Optimizations +- **Squared ReLU Activation**: Enhanced non-linearity for better gradient flow +- **QK LayerNorm**: Applies LayerNorm to query and key projections for training stability +- **RoPE**: Rotary Position Embeddings with base 10000 +- **Multi-Query Attention**: Efficient attention with shared key-value heads +- **Selective State Space**: Mamba-2 architecture with selective gating + +### Extended Context (Nano v2) +- **128K Context Window**: Nemotron Nano v2 models support up to 128K tokens +- **Efficient Long-Range Modeling**: Hybrid architecture optimized for long sequences + +## Conversion with 🤗 Hugging Face + +### Load HF → Megatron + +#### Nemotron H Models +```python +from megatron.bridge import AutoBridge + +# Example: Nemotron H 8B +bridge = AutoBridge.from_hf_pretrained("nvidia/Nemotron-H-8B-Base-8K", trust_remote_code=True) +provider = bridge.to_megatron_provider() + +# Configure parallelism before instantiating the model +provider.tensor_model_parallel_size = 2 +provider.pipeline_model_parallel_size = 1 +provider.context_parallel_size = 1 +provider.sequence_parallel = True + +provider.finalize() +model = provider.provide_distributed_model(wrap_with_ddp=False) + +# Other models: +# bridge = AutoBridge.from_hf_pretrained("nvidia/Nemotron-H-4B-Base-8K", trust_remote_code=True) +# bridge = AutoBridge.from_hf_pretrained("nvidia/Nemotron-H-47B-Base-8K", trust_remote_code=True) +# bridge = AutoBridge.from_hf_pretrained("nvidia/Nemotron-H-56B-Base-8K", trust_remote_code=True) +``` + +#### Nemotron Nano v2 Models +```python +from megatron.bridge import AutoBridge + +# Example: Nemotron Nano 9B v2 +bridge = AutoBridge.from_hf_pretrained("nvidia/NVIDIA-Nemotron-Nano-9B-v2-Base", trust_remote_code=True) +provider = bridge.to_megatron_provider() + +# Configure parallelism +provider.tensor_model_parallel_size = 2 +provider.pipeline_model_parallel_size = 1 +provider.context_parallel_size = 1 +provider.sequence_parallel = True + +provider.finalize() +model = provider.provide_distributed_model(wrap_with_ddp=False) + +# For instruct variant: +# bridge = AutoBridge.from_hf_pretrained("nvidia/NVIDIA-Nemotron-Nano-9B-v2", trust_remote_code=True) + +# For 12B model: +# bridge = AutoBridge.from_hf_pretrained("nvidia/NVIDIA-Nemotron-Nano-12B-v2-Base", trust_remote_code=True) +``` + +### Export Megatron → HF +```python +# Convert from a Megatron checkpoint directory to HF format +bridge.export_ckpt( + megatron_path="/results/nemotronh_8b/checkpoints/iter_0500000", + hf_path="./nemotronh-8b-hf-export", +) +``` + +## Examples + +- Checkpoint conversion: [examples/conversion/convert_checkpoints.py](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/conversion/convert_checkpoints.py) +- Training scripts: [examples/models/train_any_basic.py](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/models/train_any_basic.py) + +## Finetuning Recipes + +### Nemotron H 4B Finetuning + +#### LoRA Finetuning +```python +from megatron.bridge.recipes.nemotronh import nemotronh_4b_finetune_config + +cfg = nemotronh_4b_finetune_config( + tokenizer_path="nvidia/Nemotron-H-4B-Base-8K", + name="nemotronh_4b_lora", + pretrained_checkpoint="path/to/nemotronh/4b/checkpoint", + peft="lora", # or "dora" for DoRA + train_iters=1000, + global_batch_size=128, + finetune_lr=1e-4, +) +``` + +#### Full Supervised Finetuning (SFT) +```python +cfg = nemotronh_4b_finetune_config( + tokenizer_path="nvidia/Nemotron-H-4B-Base-8K", + name="nemotronh_4b_sft", + pretrained_checkpoint="path/to/nemotronh/4b/checkpoint", + peft=None, # Full supervised finetuning + train_iters=1000, + global_batch_size=128, + finetune_lr=5e-6, # Lower LR for full SFT +) +``` + +### Nemotron H 8B Finetuning + +```python +from megatron.bridge.recipes.nemotronh import nemotronh_8b_finetune_config + +# LoRA finetuning +cfg = nemotronh_8b_finetune_config( + tokenizer_path="nvidia/Nemotron-H-8B-Base-8K", + name="nemotronh_8b_lora", + pretrained_checkpoint="path/to/nemotronh/8b/checkpoint", + peft="lora", + train_iters=1000, + global_batch_size=128, + finetune_lr=1e-4, +) +``` + +### Nemotron H 47B Finetuning + +```python +from megatron.bridge.recipes.nemotronh import nemotronh_47b_finetune_config + +# LoRA finetuning (recommended for 47B) +cfg = nemotronh_47b_finetune_config( + tokenizer_path="nvidia/Nemotron-H-47B-Base-8K", + name="nemotronh_47b_lora", + pretrained_checkpoint="path/to/nemotronh/47b/checkpoint", + peft="lora", + train_iters=1000, + global_batch_size=128, + finetune_lr=1e-4, +) +``` + +### Nemotron H 56B Finetuning + +```python +from megatron.bridge.recipes.nemotronh import nemotronh_56b_finetune_config + +# LoRA finetuning (recommended for 56B) +cfg = nemotronh_56b_finetune_config( + tokenizer_path="nvidia/Nemotron-H-56B-Base-8K", + name="nemotronh_56b_lora", + pretrained_checkpoint="path/to/nemotronh/56b/checkpoint", + peft="lora", + train_iters=1000, + global_batch_size=128, + finetune_lr=1e-4, +) +``` + +### Nemotron Nano 9B v2 Finetuning + +```python +from megatron.bridge.recipes.nemotronh import nemotron_nano_9b_v2_finetune_config + +# LoRA finetuning +cfg = nemotron_nano_9b_v2_finetune_config( + tokenizer_path="nvidia/NVIDIA-Nemotron-Nano-9B-v2-Base", + name="nano_9b_v2_lora", + pretrained_checkpoint="path/to/nano/9b/v2/checkpoint", + peft="lora", + train_iters=1000, + global_batch_size=128, + seq_length=2048, # Can use up to 128K + finetune_lr=1e-4, +) +``` + +### Nemotron Nano 12B v2 Finetuning + +```python +from megatron.bridge.recipes.nemotronh import nemotron_nano_12b_v2_finetune_config + +# LoRA finetuning +cfg = nemotron_nano_12b_v2_finetune_config( + tokenizer_path="nvidia/NVIDIA-Nemotron-Nano-12B-v2-Base", + name="nano_12b_v2_lora", + pretrained_checkpoint="path/to/nano/12b/v2/checkpoint", + peft="lora", + train_iters=1000, + global_batch_size=128, + seq_length=2048, # Can use up to 128K + finetune_lr=1e-4, +) +``` + +## Default Configurations + +### Nemotron H Models + +#### 4B - LoRA (1 node, 8 GPUs) +- TP=1, PP=1, CP=1, LR=1e-4 +- Sequence Parallel: False +- Precision: BF16 mixed +- Optimized for single-GPU finetuning + +#### 4B - Full SFT (1 node, 8 GPUs) +- TP=1, PP=1, CP=1, LR=5e-6 +- Sequence Parallel: False +- Precision: BF16 mixed + +#### 8B - LoRA (1 node, 8 GPUs) +- TP=1, PP=1, CP=1, LR=1e-4 +- Sequence Parallel: False +- Precision: BF16 mixed + +#### 8B - Full SFT (1 node, 8 GPUs) +- TP=2, PP=1, CP=1, LR=5e-6 +- Sequence Parallel: True +- Precision: BF16 mixed + +#### 47B - LoRA (2+ nodes) +- TP=4, PP=1, CP=1, LR=1e-4 +- Sequence Parallel: False +- Precision: FP8 hybrid (recommended) + +#### 47B - Full SFT (4+ nodes) +- TP=8, PP=1, CP=1, LR=5e-6 +- Sequence Parallel: True +- Precision: FP8 hybrid + +#### 56B - LoRA (2+ nodes) +- TP=4, PP=1, CP=1, LR=1e-4 +- Sequence Parallel: False +- Precision: FP8 hybrid (recommended) + +#### 56B - Full SFT (4+ nodes) +- TP=8, PP=1, CP=1, LR=5e-6 +- Sequence Parallel: True +- Precision: FP8 hybrid + +### Nemotron Nano v2 Models + +#### 9B - LoRA (1 node, 8 GPUs) +- TP=2, PP=1, CP=1, LR=1e-4 +- Sequence Parallel: True +- Precision: BF16 mixed +- Context: Up to 128K tokens + +#### 9B - Full SFT (1 node, 8 GPUs) +- TP=2, PP=1, CP=1, LR=1e-4 +- Sequence Parallel: True +- Precision: BF16 mixed + +#### 12B - LoRA (2 nodes, 16 GPUs) +- TP=4, PP=1, CP=1, LR=1e-4 +- Sequence Parallel: True +- Precision: FP8 hybrid (recommended) +- Context: Up to 128K tokens + +#### 12B - Full SFT (2 nodes, 16 GPUs) +- TP=4, PP=1, CP=1, LR=1e-4 +- Sequence Parallel: True +- Precision: FP8 hybrid + +## API Reference + +### Nemotron H +- Nemotron H recipes: [bridge.recipes.nemotronh](/../../apidocs/bridge/bridge.recipes.nemotronh) +- Nemotron H model providers: [bridge.models.nemotronh](/../../apidocs/bridge/bridge.models.nemotronh) + +### Nemotron Nano v2 +- Nemotron Nano v2 recipes: [bridge.recipes.nemotronh.nemotron_nano_v2](/../../apidocs/bridge/bridge.recipes.nemotronh) +- Nemotron Nano v2 model providers: [bridge.models.nemotronh.NemotronNanoModelProvider9Bv2](/../../apidocs/bridge/bridge.models.nemotronh) + +## Performance Optimizations + +### Memory Efficiency +- **Selective Recomputation**: Reduces activation memory for larger models +- **Sequence Parallelism**: Distributes sequence dimension across GPUs (enabled for 8B+) +- **Context Parallelism**: Support for ultra-long sequences (Nano v2) +- **Manual GC**: Aggressive garbage collection for stable memory usage +- **Precision-aware optimizer**: BF16/FP8 gradients with FP32 master weights + +### Compute Efficiency +- **Mamba-2 Optimizations**: Efficient selective state space computations +- **Hybrid Architecture**: Balanced mix of Mamba and Attention layers +- **Squared ReLU**: Efficient activation function with good gradient properties +- **RoPE Fusion**: Optional optimization for position embeddings +- **Multi-Query Attention**: Reduced KV cache memory and compute + +### Hybrid Pattern Optimization +The hybrid override pattern determines which layers use Mamba (M) vs Attention (*): +- **Mamba layers**: Fast, memory-efficient, good for long-range dependencies +- **Attention layers**: Better for complex reasoning and multi-token relationships +- **Optimal patterns**: Pre-configured per model size based on extensive experimentation + +## Pipeline Parallelism Layouts + +Nemotron H models support several PP configurations with pre-defined layouts: +- **PP=1**: No pipelining (default for most configurations) +- **PP=2**: Supported with symmetric layer splits +- **PP=4**: Supported for larger models (47B, 56B) +- **VP (Virtual Pipeline)**: Supported for reducing pipeline bubbles + +## Hugging Face Model Cards + +### Nemotron H Models +- **4B Base**: [nvidia/Nemotron-H-4B-Base-8K](https://huggingface.co/nvidia/Nemotron-H-4B-Base-8K) +- **8B Base**: [nvidia/Nemotron-H-8B-Base-8K](https://huggingface.co/nvidia/Nemotron-H-8B-Base-8K) +- **47B Base**: [nvidia/Nemotron-H-47B-Base-8K](https://huggingface.co/nvidia/Nemotron-H-47B-Base-8K) +- **56B Base**: [nvidia/Nemotron-H-56B-Base-8K](https://huggingface.co/nvidia/Nemotron-H-56B-Base-8K) + +### Nemotron Nano v2 Models +- **9B Base**: [nvidia/NVIDIA-Nemotron-Nano-9B-v2-Base](https://huggingface.co/nvidia/NVIDIA-Nemotron-Nano-9B-v2-Base) +- **9B Instruct**: [nvidia/NVIDIA-Nemotron-Nano-9B-v2](https://huggingface.co/nvidia/NVIDIA-Nemotron-Nano-9B-v2) +- **12B Base**: [nvidia/NVIDIA-Nemotron-Nano-12B-v2-Base](https://huggingface.co/nvidia/NVIDIA-Nemotron-Nano-12B-v2-Base) +- **12B Instruct**: [nvidia/NVIDIA-Nemotron-Nano-12B-v2](https://huggingface.co/nvidia/NVIDIA-Nemotron-Nano-12B-v2) + +## Technical Resources + +### Research Papers +- **Nemotron Technical Report**: [arXiv:2508.14444](https://arxiv.org/abs/2508.14444) +- **Mamba-2**: [Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality](https://arxiv.org/abs/2405.21060) + +## Related Documentation + +- Recipe usage and customization: [Recipe usage](/../../recipe-usage) +- Training configuration: [Configuration overview](/../../training/config-container-overview) +- Training entry points: [Entry points](/../../training/entry-points) +- PEFT methods (LoRA, DoRA): [PEFT Guide](/../../training/peft) diff --git a/fern/v0.2.0/pages/models/llm/olmoe.mdx b/fern/v0.2.0/pages/models/llm/olmoe.mdx new file mode 100644 index 0000000000..80c05abbb6 --- /dev/null +++ b/fern/v0.2.0/pages/models/llm/olmoe.mdx @@ -0,0 +1,183 @@ +--- +title: "OLMoE" +description: "" +--- + +[OLMoE](https://huggingface.co/allenai/OLMoE-1B-7B-0125) is a 7B-parameter Mixture-of-Experts (MoE) model from **Allen Institute for AI (AI2)** featuring 64 experts with top-8 routing. The model is designed to be fully open-source, with training data, code, and model weights publicly available. It's named "OLMoE-1B-7B" where 1B refers to the activated parameters and 7B refers to the total parameters. + +The latest version (OLMoE-1B-7B-0125, released January 2025) is an improved version of the original September 2024 release (OLMoE-1B-7B-0924), trained on 5T tokens with performance improvements across multiple benchmarks. + +The model features 16 decoder layers with 64 routed experts per layer, activating 8 experts per token for a total of approximately 1.3B active parameters per forward pass out of 7B total. + +OLMoE models are supported via the Bridge system with specialized configurations for MoE optimizations. + +## Model Architecture + +- **Parameters**: 7B total, 1.3B activated per forward pass +- **Layers**: 16 decoder layers +- **Attention**: Multi-query attention with QK LayerNorm and RoPE +- **MoE**: 64 routed experts per layer with top-8 routing +- **Hidden size**: 2048 +- **FFN hidden size**: 1024 (dense layers), 1024 (expert layers) +- **Attention heads**: 16 query heads, 16 key-value groups +- **Vocab size**: 50,304 +- **Context Length**: 4K tokens +- **Activation**: SiLU with gated linear units +- **Training**: 5T tokens (OLMoE-1B-7B-0125) + +## Key Features + +- **QK LayerNorm**: Applies LayerNorm to query and key projections for training stability +- **RoPE**: Rotary Position Embeddings with base 10000 +- **MoE Routing**: Softmax-based router with auxiliary loss for load balancing +- **Router Pre-Softmax**: Pre-softmax routing scores +- **Grouped GEMM**: Optimized grouped matrix multiplications for expert computation + +## Conversion with 🤗 Hugging Face + +### Load HF → Megatron +```python +from megatron.bridge import AutoBridge + +# Example: OLMoE-1B-7B-0125 (latest version) +bridge = AutoBridge.from_hf_pretrained("allenai/OLMoE-1B-7B-0125") +provider = bridge.to_megatron_provider() + +# Configure parallelism before instantiating the model +provider.tensor_model_parallel_size = 1 +provider.pipeline_model_parallel_size = 1 +provider.expert_model_parallel_size = 8 +provider.sequence_parallel = False + +provider.finalize() +model = provider.provide_distributed_model(wrap_with_ddp=False) +# You can also use older versions: +# bridge = AutoBridge.from_hf_pretrained("allenai/OLMoE-1B-7B-0924") +``` + +### Export Megatron → HF +```python +# Convert from a Megatron checkpoint directory to HF format +bridge.export_ckpt( + megatron_path="/results/olmoe_7b/checkpoints/iter_0500000", + hf_path="./olmoe-hf-export", +) +``` + +## Examples + +- Checkpoint conversion: [examples/conversion/convert_checkpoints.py](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/conversion/convert_checkpoints.py) + +## Recipes + +See: [bridge.recipes.olmoe](/../../apidocs/bridge/bridge.recipes.olmoe) + +### Available Recipes + +- **Pretrain recipes**: + - `olmoe_7b_pretrain_config`: Pre-training for OLMoE-7B (7B parameters, 1.3B activated per token) + +- **Finetune recipes**: + - `olmoe_7b_finetune_config`: Finetuning for OLMoE-7B with PEFT support (LoRA, DoRA) + +### Parallelism Configurations + +| Model | Mode | TP | PP | EP | Total GPUs | Use Case | +|-------|------|----|----|----|-----------:|----------| +| **OLMoE-7B** | Pretrain | 1 | 1 | 8 | 8 | Pre-training (single node) | +| **OLMoE-7B** | Full SFT | 1 | 1 | 8 | 8 | Full supervised finetuning | +| **OLMoE-7B** | LoRA/DoRA | 1 | 1 | 1 | 8 | PEFT finetuning (single node) | + +**Key Features**: +- **Expert Parallelism**: EP=8 for efficient MoE training (64 experts) +- **Selective Recomputation**: Enabled by default for memory optimization +- **RoPE Fusion**: Optional optimization for MLA (`apply_rope_fusion=True`) +- **MoE Optimizations**: Grouped GEMM and permute fusion enabled by default + +**Performance Optimizations**: +- **MoE Permute Fusion**: Fused expert permutation operations +- **Grouped GEMM**: Optimized expert computation +- **Router Load Balancing**: Auxiliary loss for balanced expert utilization +- **Manual GC**: Aggressive garbage collection (interval=5) +- **Precision-Aware Optimizer**: BF16 gradients and optimizer states with FP32 master weights + +**Pipeline Layouts** (optional): +- **PP=1**: No pipelining (default) +- **PP=2**: 8+8 layer split with embedding/loss +- **PP=4**: 4+4+4+4 layer split +- **VP**: PP=2,VP=2 supported + +### Pre-training Example + +```python +from megatron.bridge.recipes.olmoe import olmoe_7b_pretrain_config + +cfg = olmoe_7b_pretrain_config( + name="olmoe_pretrain", + data_paths=["/path/to/dataset.nvjsonl"], + dir="/results/olmoe_7b", + train_iters=500_000, + global_batch_size=2048, + seq_length=4096, + # Uses TP=1, PP=1, EP=8 (8 GPUs) automatically +) +``` + +### Finetuning Examples + +#### Full Finetuning + +```python +from megatron.bridge.recipes.olmoe import olmoe_7b_finetune_config + +cfg = olmoe_7b_finetune_config( + tokenizer_path="allenai/OLMoE-1B-7B-0125", + name="olmoe_full_sft", + pretrained_checkpoint="path/to/olmoe/checkpoint", + peft=None, # Full supervised finetuning + train_iters=1000, + global_batch_size=128, + finetune_lr=5e-6, + # Uses TP=1, PP=1, EP=8 (8 GPUs) automatically +) +``` + +#### LoRA Finetuning + +```python +from megatron.bridge.recipes.olmoe import olmoe_7b_finetune_config + +cfg = olmoe_7b_finetune_config( + tokenizer_path="allenai/OLMoE-1B-7B-0125", + name="olmoe_lora_finetune", + pretrained_checkpoint="path/to/olmoe/checkpoint", + peft="lora", # or "dora" for DoRA + train_iters=1000, + global_batch_size=128, + finetune_lr=1e-4, + # Uses TP=1, PP=1, EP=1 (8 GPUs) automatically +) +``` + +## Hugging Face model cards + +### Latest (January 2025) +- OLMoE-1B-7B-0125 (Base): [allenai/OLMoE-1B-7B-0125](https://huggingface.co/allenai/OLMoE-1B-7B-0125) +- OLMoE-1B-7B-0125-SFT: [allenai/OLMoE-1B-7B-0125-SFT](https://huggingface.co/allenai/OLMoE-1B-7B-0125-SFT) +- OLMoE-1B-7B-0125-Instruct: [allenai/OLMoE-1B-7B-0125-Instruct](https://huggingface.co/allenai/OLMoE-1B-7B-0125-Instruct) + +### Previous (September 2024) +- OLMoE-1B-7B-0924 (Base): [allenai/OLMoE-1B-7B-0924](https://huggingface.co/allenai/OLMoE-1B-7B-0924) +- OLMoE-1B-7B-0924-Instruct: [allenai/OLMoE-1B-7B-0924-Instruct](https://huggingface.co/allenai/OLMoE-1B-7B-0924-Instruct) + +## Technical Resources + +- OLMoE Paper: [OLMoE: Open Mixture-of-Experts Language Models](https://arxiv.org/abs/2409.02060) +- OLMoE Model Card (Latest): [HuggingFace Model Card](https://huggingface.co/allenai/OLMoE-1B-7B-0125) +- OLMoE GitHub Repository: [allenai/OLMoE](https://github.com/allenai/OLMoE) + +## Related docs + +- Recipe usage and customization: [Recipe usage](/../../recipe-usage) +- Training configuration: [Configuration overview](/../../training/config-container-overview) +- Training entry points: [Entry points](/../../training/entry-points) diff --git a/fern/v0.2.0/pages/models/llm/qwen.mdx b/fern/v0.2.0/pages/models/llm/qwen.mdx new file mode 100644 index 0000000000..073a174b2f --- /dev/null +++ b/fern/v0.2.0/pages/models/llm/qwen.mdx @@ -0,0 +1,574 @@ +--- +title: "Qwen" +description: "" +--- + +[Qwen](https://huggingface.co/Qwen) is a family of large language models developed by Alibaba Cloud, including dense models (Qwen2, Qwen2.5, Qwen3) and Mixture-of-Experts models (Qwen3 MoE, Qwen3-Next). The models feature innovations like QK layernorm, Gated-Delta Networks, and Zero-Centered RMSNorm for improved training stability and performance. + +Qwen family models are supported via the Bridge system with auto-detected configuration and weight mapping. + +## Available Models + +Megatron Bridge supports the following Qwen model variants: + +### Dense Models +- **Qwen2**: 0.5B, 1.5B, 7B, 72B +- **Qwen2.5**: 0.5B, 1.5B, 7B, 14B, 32B, 72B +- **Qwen3**: 0.6B, 1.7B, 4B, 8B, 14B, 32B + +### MoE Models +- **Qwen3 MoE**: 30B (3B activated), 235B (22B activated) +- **Qwen3-Next**: 80B (3B activated) + +## Model Architecture Features + +### Common Features +- **Pre-normalization**: RMSNorm before each transformer sub-layer +- **SwiGLU Activation**: Gated linear units in the feedforward network +- **Rotary Positional Embeddings (RoPE)**: Relative position encoding +- **Grouped Query Attention (GQA)**: Memory-efficient attention mechanism + +### Qwen3-Next-Specific Features +- **Gated-Delta Networks**: Advanced gating mechanism for improved learning +- **Zero-Centered RMSNorm**: Centered normalization for training stability +- **Multi-Token Prediction (MTP)**: Auxiliary training objective + +### Qwen3-Specific Features +- **QK Layernorm**: Layer normalization on query and key projections +- **QK Layernorm Weight Decay**: Weight decay applied during training + +### Qwen2-Specific Features +- **Bias in QKV**: Bias terms in query, key, value projections + +--- + +## Qwen3-Next + +### Conversion with 🤗 Hugging Face + +#### Load HF → Megatron + +```python +from megatron.bridge import AutoBridge + +# Example: Qwen3-Next-80B-A3B +bridge = AutoBridge.from_hf_pretrained("Qwen/Qwen3-Next-80B-A3B-Instruct") +provider = bridge.to_megatron_provider() + +# Optionally configure parallelism before instantiating the model +provider.tensor_model_parallel_size = 2 +provider.pipeline_model_parallel_size = 8 +provider.expert_model_parallel_size = 16 + +model = provider.provide_distributed_model(wrap_with_ddp=False) +``` + +#### Import Checkpoint from HF + +```bash +python examples/conversion/convert_checkpoints.py import \ + --hf-model Qwen/Qwen3-Next-80B-A3B-Instruct \ + --megatron-path /checkpoints/qwen3_next_80b_megatron +``` + +#### Export Megatron → HF + +```python +from megatron.bridge import AutoBridge + +# Load the bridge from HF model ID +bridge = AutoBridge.from_hf_pretrained("Qwen/Qwen3-Next-80B-A3B-Instruct") + +# Export a trained Megatron checkpoint to HF format +bridge.export_ckpt( + megatron_path="/results/qwen3_next_80b/checkpoints/iter_0000500", + hf_path="/exports/qwen3_next_80b_hf", +) +``` + +#### Run Inference on Converted Checkpoint + +```bash +python examples/conversion/hf_to_megatron_generate_text.py \ + --hf_model_path Qwen/Qwen3-Next-80B-A3B-Instruct \ + --megatron_model_path /checkpoints/qwen3_next_80b_megatron \ + --prompt "What is artificial intelligence?" \ + --max_new_tokens 100 \ + --tp 2 \ + --pp 8 \ + --ep 16 +``` + +### Recipes + +#### Available Recipes +- `qwen3_next_80b_a3b_pretrain_config`: Pre-training for Qwen3-Next-80B-A3B +- `qwen3_next_80b_a3b_finetune_config`: Finetuning for Qwen3-Next-80B-A3B (Full SFT only) + +#### Parallelism Configuration + +| Model | Mode | TP | PP | EP | Total GPUs | Use Case | +|-------|------|----|----|----|-----------:|----------| +| **Qwen3-Next-80B** | Pretrain | 2 | 8 | 16 | 256 | Pre-training (32 nodes) | +| **Qwen3-Next-80B** | Full SFT | 2 | 8 | 16 | 256 | Full supervised finetuning (32 nodes) | + +#### Pre-training Example + +```python +from megatron.bridge.recipes.qwen import qwen3_next_80b_a3b_pretrain_config + +config = qwen3_next_80b_a3b_pretrain_config( + name="qwen3_next_80b_pretrain", + data_paths=["/path/to/dataset.nvjsonl"], + dir="/results/qwen3_next_80b", + train_iters=500_000, + global_batch_size=2048, + seq_length=4096, + # Uses TP=2, PP=8, EP=16 (256 GPUs) automatically +) +``` + +#### Finetuning Example + +```python +from megatron.bridge.recipes.qwen import qwen3_next_80b_a3b_finetune_config + +config = qwen3_next_80b_a3b_finetune_config( + name="qwen3_next_80b_full_sft", + pretrained_checkpoint="/results/qwen3_next_80b/checkpoints/iter_0500000", + peft=None, # Full supervised finetuning + train_iters=1000, + global_batch_size=64, + finetune_lr=5e-6, + # Uses TP=2, PP=8, EP=16 (256 GPUs) automatically +) +``` + +**Note**: PEFT (LoRA/DoRA) finetuning is not currently available for Qwen3-Next models. + +### Hugging Face Model Cards + +- Qwen3-Next-80B-A3B-Instruct: https://huggingface.co/Qwen/Qwen3-Next-80B-A3B-Instruct +- Qwen3-Next-80B-A3B-Thinking: https://huggingface.co/Qwen/Qwen3-Next-80B-A3B-Thinking + +--- + +## Qwen3 MoE + +### Conversion with 🤗 Hugging Face + +#### Load HF → Megatron + +```python +from megatron.bridge import AutoBridge + +# Example: Qwen3-30B-A3B +bridge = AutoBridge.from_hf_pretrained("Qwen/Qwen3-30B-A3B") +provider = bridge.to_megatron_provider() + +# Optionally configure parallelism before instantiating the model +provider.tensor_model_parallel_size = 1 +provider.pipeline_model_parallel_size = 1 +provider.expert_model_parallel_size = 8 + +model = provider.provide_distributed_model(wrap_with_ddp=False) +``` + +#### Import Checkpoint from HF + +```bash +python examples/conversion/convert_checkpoints.py import \ + --hf-model Qwen/Qwen3-30B-A3B \ + --megatron-path /checkpoints/qwen3_30b_a3b_megatron +``` + +#### Export Megatron → HF + +```python +from megatron.bridge import AutoBridge + +# Load the bridge from HF model ID +bridge = AutoBridge.from_hf_pretrained("Qwen/Qwen3-30B-A3B") + +# Export a trained Megatron checkpoint to HF format +bridge.export_ckpt( + megatron_path="/results/qwen3_30b_a3b/checkpoints/iter_0000500", + hf_path="/exports/qwen3_30b_a3b_hf", +) +``` + +#### Run Inference on Converted Checkpoint + +```bash +python examples/conversion/hf_to_megatron_generate_text.py \ + --hf_model_path Qwen/Qwen3-30B-A3B \ + --megatron_model_path /checkpoints/qwen3_30b_a3b_megatron \ + --prompt "What is artificial intelligence?" \ + --max_new_tokens 100 \ + --ep 8 +``` + +### Recipes + +#### Available Recipes +- `qwen3_30b_a3b_pretrain_config`: Pre-training for Qwen3-30B-A3B (30B parameters, 3B activated) +- `qwen3_235b_a22b_pretrain_config`: Pre-training for Qwen3-235B-A22B (235B parameters, 22B activated) +- `qwen3_30b_a3b_finetune_config`: Finetuning for Qwen3-30B-A3B with PEFT support +- `qwen3_235b_a22b_finetune_config`: Finetuning for Qwen3-235B-A22B with PEFT support + +#### Parallelism Configuration + +| Model | Mode | TP | PP | EP | Total GPUs | Use Case | +|-------|------|----|----|----|-----------:|----------| +| **Qwen3-30B-A3B** | Pretrain | 1 | 1 | 8 | 8 | Pre-training (single node) | +| **Qwen3-30B-A3B** | Full SFT | 1 | 1 | 8 | 8 | Full supervised finetuning | +| **Qwen3-30B-A3B** | LoRA/DoRA | 1 | 1 | 8 | 8 | PEFT finetuning (single node) | +| **Qwen3-235B-A22B** | Pretrain | 2 | 8 | 32 | 512 | Pre-training (64 nodes) | +| **Qwen3-235B-A22B** | Full SFT | 2 | 8 | 32 | 512 | Full supervised finetuning (64 nodes) | +| **Qwen3-235B-A22B** | LoRA/DoRA | 2 | 8 | 32 | 512 | PEFT finetuning (64 nodes) | + +#### Pre-training Examples + +**Qwen3-30B-A3B:** + +```python +from megatron.bridge.recipes.qwen import qwen3_30b_a3b_pretrain_config + +config = qwen3_30b_a3b_pretrain_config( + name="qwen3_30b_a3b_pretrain", + data_paths=["/path/to/dataset.nvjsonl"], + dir="/results/qwen3_30b_a3b", + train_iters=500_000, + global_batch_size=2048, + seq_length=4096, + # Uses TP=1, PP=1, EP=8 (8 GPUs) automatically +) +``` + +**Qwen3-235B-A22B** + +```python +from megatron.bridge.recipes.qwen import qwen3_235b_a22b_pretrain_config + +config = qwen3_235b_a22b_pretrain_config( + name="qwen3_235b_a22b_pretrain", + data_paths=["/path/to/dataset.nvjsonl"], + dir="/results/qwen3_235b_a22b", + train_iters=500_000, + global_batch_size=4096, + seq_length=4096, + # Uses TP=2, PP=8, EP=32 (512 GPUs) automatically +) +``` + +#### Finetuning Examples + +**Full Finetuning (30B):** + +```python +from megatron.bridge.recipes.qwen import qwen3_30b_a3b_finetune_config + +config = qwen3_30b_a3b_finetune_config( + name="qwen3_30b_a3b_full_sft", + pretrained_checkpoint="/results/qwen3_30b_a3b/checkpoints/iter_0500000", + peft=None, + train_iters=1000, + global_batch_size=64, + finetune_lr=5e-6, + # Uses TP=1, PP=1, EP=8 (8 GPUs) automatically +) +``` + +**LoRA Finetuning (30B):** + +```python +from megatron.bridge.recipes.qwen import qwen3_30b_a3b_finetune_config + +config = qwen3_30b_a3b_finetune_config( + name="qwen3_30b_a3b_lora", + pretrained_checkpoint="/results/qwen3_30b_a3b/checkpoints/iter_0500000", + peft="lora", # or "dora" + train_iters=1000, + global_batch_size=128, + finetune_lr=1e-4, + # Uses TP=1, PP=1, EP=8 (8 GPUs) automatically +) +``` + +### Hugging Face Model Cards + +- Qwen3-30B-A3B: https://huggingface.co/Qwen/Qwen3-30B-A3B +- Qwen3-235B-A22B: https://huggingface.co/Qwen/Qwen3-235B-A22B + +--- + +## Qwen3 + +### Conversion with 🤗 Hugging Face + +#### Load HF → Megatron + +```python +from megatron.bridge import AutoBridge + +# Example: Qwen3-8B +bridge = AutoBridge.from_hf_pretrained("Qwen/Qwen3-8B") +provider = bridge.to_megatron_provider() + +# Optionally configure parallelism before instantiating the model +provider.tensor_model_parallel_size = 2 +provider.pipeline_model_parallel_size = 1 + +model = provider.provide_distributed_model(wrap_with_ddp=False) +``` + +#### Import Checkpoint from HF + +```bash +python examples/conversion/convert_checkpoints.py import \ + --hf-model Qwen/Qwen3-8B \ + --megatron-path /checkpoints/qwen3_8b_megatron +``` + +#### Export Megatron → HF + +```python +from megatron.bridge import AutoBridge + +# Load the bridge from HF model ID +bridge = AutoBridge.from_hf_pretrained("Qwen/Qwen3-8B") + +# Export a trained Megatron checkpoint to HF format +bridge.export_ckpt( + megatron_path="/results/qwen3_8b/checkpoints/iter_0000500", + hf_path="/exports/qwen3_8b_hf", +) +``` + +#### Run Inference on Converted Checkpoint + +```bash +python examples/conversion/hf_to_megatron_generate_text.py \ + --hf_model_path Qwen/Qwen3-8B \ + --megatron_model_path /checkpoints/qwen3_8b_megatron \ + --prompt "What is artificial intelligence?" \ + --max_new_tokens 100 \ + --tp 2 +``` + +### Recipes + +#### Available Recipes +- **Pretrain recipes**: `qwen3_600m_pretrain_config`, `qwen3_1p7b_pretrain_config`, `qwen3_4b_pretrain_config`, `qwen3_8b_pretrain_config`, `qwen3_14b_pretrain_config`, `qwen3_32b_pretrain_config` +- **Finetune recipes**: `qwen3_600m_finetune_config`, `qwen3_1p7b_finetune_config`, `qwen3_4b_finetune_config`, `qwen3_8b_finetune_config`, `qwen3_14b_finetune_config`, `qwen3_32b_finetune_config` (all with PEFT support) + +#### Parallelism Configuration + +| Model | Mode | TP | PP | Total GPUs | Use Case | +|-------|------|----|----|------------|----------| +| **Qwen3 (0.6B-4B)** | Pretrain | 1 | 1 | 8 | Pre-training (single node) | +| **Qwen3 (0.6B-4B)** | Full SFT | 1 | 1 | 8 | Full supervised finetuning | +| **Qwen3 (0.6B-4B)** | LoRA/DoRA | 1 | 1 | 8 | PEFT finetuning (single node) | +| **Qwen3 (8B-14B)** | Pretrain | 2 | 1 | 16 | Pre-training (2 nodes) | +| **Qwen3 (8B-14B)** | Full SFT | 2 | 1 | 16 | Full supervised finetuning (2 nodes) | +| **Qwen3 (8B-14B)** | LoRA/DoRA | 1 | 1 | 8 | PEFT finetuning (single node!) | +| **Qwen3-32B** | Pretrain | 4 | 1 | 32 | Pre-training (4 nodes) | +| **Qwen3-32B** | Full SFT | 4 | 1 | 32 | Full supervised finetuning (4 nodes) | +| **Qwen3-32B** | LoRA/DoRA | 2 | 1 | 16 | PEFT finetuning (2 nodes) | + +#### Pre-training Example + +```python +from megatron.bridge.recipes.qwen import qwen3_8b_pretrain_config + +config = qwen3_8b_pretrain_config( + name="qwen3_8b_pretrain", + data_paths=["/path/to/dataset.nvjsonl"], + dir="/results/qwen3_8b", + train_iters=500_000, + global_batch_size=2048, + seq_length=4096, + # Uses TP=2, PP=1 (16 GPUs) automatically +) +``` + +#### Finetuning Examples + +**Full Finetuning (8B):** + +```python +from megatron.bridge.recipes.qwen import qwen3_8b_finetune_config + +config = qwen3_8b_finetune_config( + name="qwen3_8b_full_sft", + pretrained_checkpoint="/results/qwen3_8b/checkpoints/iter_0500000", + peft=None, + train_iters=1000, + global_batch_size=64, + finetune_lr=5e-6, + # Uses TP=2, PP=1 (16 GPUs) automatically +) +``` + +**LoRA Finetuning (8B):** + +```python +from megatron.bridge.recipes.qwen import qwen3_8b_finetune_config + +config = qwen3_8b_finetune_config( + name="qwen3_8b_lora", + pretrained_checkpoint="/results/qwen3_8b/checkpoints/iter_0500000", + peft="lora", # or "dora" + train_iters=1000, + global_batch_size=128, + finetune_lr=1e-4, + # Uses TP=1, PP=1 (8 GPUs) automatically +) +``` + +### Hugging Face Model Cards + +- Qwen3 Collection: https://huggingface.co/collections/Qwen/qwen3 + +--- + +## Qwen2 / Qwen2.5 + +### Conversion with 🤗 Hugging Face + +#### Load HF → Megatron + +```python +from megatron.bridge import AutoBridge + +# Example: Qwen2.5-7B +bridge = AutoBridge.from_hf_pretrained("Qwen/Qwen2.5-7B") +provider = bridge.to_megatron_provider() + +# Optionally configure parallelism before instantiating the model +provider.tensor_model_parallel_size = 2 +provider.pipeline_model_parallel_size = 1 + +model = provider.provide_distributed_model(wrap_with_ddp=False) +``` + +#### Import Checkpoint from HF + +```bash +python examples/conversion/convert_checkpoints.py import \ + --hf-model Qwen/Qwen2.5-7B \ + --megatron-path /checkpoints/qwen25_7b_megatron +``` + +#### Export Megatron → HF + +```python +from megatron.bridge import AutoBridge + +# Load the bridge from HF model ID +bridge = AutoBridge.from_hf_pretrained("Qwen/Qwen2.5-7B") + +# Export a trained Megatron checkpoint to HF format +bridge.export_ckpt( + megatron_path="/results/qwen25_7b/checkpoints/iter_0000500", + hf_path="/exports/qwen25_7b_hf", +) +``` + +#### Run Inference on Converted Checkpoint + +```bash +python examples/conversion/hf_to_megatron_generate_text.py \ + --hf_model_path Qwen/Qwen2.5-7B \ + --megatron_model_path /checkpoints/qwen25_7b_megatron \ + --prompt "What is artificial intelligence?" \ + --max_new_tokens 100 \ + --tp 2 +``` + +### Recipes + +#### Available Recipes +- **Qwen2 Pretrain**: `qwen2_500m_pretrain_config`, `qwen2_1p5b_pretrain_config`, `qwen2_7b_pretrain_config`, `qwen2_72b_pretrain_config` +- **Qwen2.5 Pretrain**: `qwen25_500m_pretrain_config`, `qwen25_1p5b_pretrain_config`, `qwen25_7b_pretrain_config`, `qwen25_14b_pretrain_config`, `qwen25_32b_pretrain_config`, `qwen25_72b_pretrain_config` +- **Qwen2 Finetune**: `qwen2_500m_finetune_config`, `qwen2_1p5b_finetune_config`, `qwen2_7b_finetune_config`, `qwen2_72b_finetune_config` (all with PEFT support) +- **Qwen2.5 Finetune**: `qwen25_500m_finetune_config`, `qwen25_1p5b_finetune_config`, `qwen25_7b_finetune_config`, `qwen25_14b_finetune_config`, `qwen25_32b_finetune_config`, `qwen25_72b_finetune_config` (all with PEFT support) + +#### Parallelism Configuration + +| Model | Mode | TP | PP | Total GPUs | Use Case | +|-------|------|----|----|------------|----------| +| **Qwen2/2.5 (0.5B-1.5B)** | Pretrain | 1 | 1 | 8 | Pre-training (single node) | +| **Qwen2/2.5 (0.5B-1.5B)** | Full SFT | 1 | 1 | 8 | Full supervised finetuning | +| **Qwen2/2.5 (0.5B-1.5B)** | LoRA/DoRA | 1 | 1 | 8 | PEFT finetuning (single node) | +| **Qwen2/2.5 (7B-14B)** | Pretrain | 2 | 1 | 16 | Pre-training (2 nodes) | +| **Qwen2/2.5 (7B-14B)** | Full SFT | 2 | 1 | 16 | Full supervised finetuning (2 nodes) | +| **Qwen2/2.5 (7B-14B)** | LoRA/DoRA | 1 | 1 | 8 | PEFT finetuning (single node!) | +| **Qwen2.5-32B** | Pretrain | 4 | 1 | 32 | Pre-training (4 nodes) | +| **Qwen2.5-32B** | Full SFT | 4 | 1 | 32 | Full supervised finetuning (4 nodes) | +| **Qwen2.5-32B** | LoRA/DoRA | 2 | 1 | 16 | PEFT finetuning (2 nodes) | +| **Qwen2/2.5-72B** | Pretrain | 8 | 1 | 64 | Pre-training (8 nodes) | +| **Qwen2/2.5-72B** | Full SFT | 8 | 1 | 64 | Full supervised finetuning (8 nodes) | +| **Qwen2/2.5-72B** | LoRA/DoRA | 4 | 1 | 32 | PEFT finetuning (4 nodes) | + +#### Pre-training Example + +```python +from megatron.bridge.recipes.qwen import qwen25_7b_pretrain_config + +config = qwen25_7b_pretrain_config( + name="qwen25_7b_pretrain", + data_paths=["/path/to/dataset.nvjsonl"], + dir="/results/qwen25_7b", + train_iters=500_000, + global_batch_size=2048, + seq_length=4096, + # Uses TP=2, PP=1 (16 GPUs) automatically +) +``` + +#### Finetuning Examples + +**Full Finetuning (7B):** + +```python +from megatron.bridge.recipes.qwen import qwen25_7b_finetune_config + +config = qwen25_7b_finetune_config( + name="qwen25_7b_full_sft", + pretrained_checkpoint="/results/qwen25_7b/checkpoints/iter_0500000", + peft=None, + train_iters=1000, + global_batch_size=64, + finetune_lr=5e-6, + # Uses TP=2, PP=1 (16 GPUs) automatically +) +``` + +**LoRA Finetuning (7B):** + +```python +from megatron.bridge.recipes.qwen import qwen25_7b_finetune_config + +config = qwen25_7b_finetune_config( + name="qwen25_7b_lora", + pretrained_checkpoint="/results/qwen25_7b/checkpoints/iter_0500000", + peft="lora", # or "dora" + train_iters=1000, + global_batch_size=128, + finetune_lr=1e-4, + # Uses TP=1, PP=1 (8 GPUs) automatically +) +``` + +### Hugging Face Model Cards + +- Qwen2 Collection: https://huggingface.co/collections/Qwen/qwen2 +- Qwen2.5 Collection: https://huggingface.co/collections/Qwen/qwen25 + +--- + +## Related Docs +- Recipe usage: [Recipe usage](/../../recipe-usage) +- Customizing the training recipe configuration: [Configuration overview](/../../training/config-container-overview) +- Training entry points: [Entry points](/../../training/entry-points) diff --git a/fern/v0.2.0/pages/models/vlm/README.mdx b/fern/v0.2.0/pages/models/vlm/README.mdx new file mode 100644 index 0000000000..d4072c48b5 --- /dev/null +++ b/fern/v0.2.0/pages/models/vlm/README.mdx @@ -0,0 +1,58 @@ +--- +title: "Vision Language Models" +description: "" +--- + +This directory contains documentation for Vision Language Models (VLMs) supported by Megatron Bridge. These models combine vision and language capabilities for multimodal AI applications. + +## Available Models + +Megatron Bridge supports the following VLM families: + +| Model | Documentation | Description | +|-------|---------------|-------------| +| **Gemma 3 VL** | [gemma3-vl.md](/gemma3-vl) | Google Gemma 3 Vision Language model | +| **Ministral 3** | [ministral3.md](/ministral3) | Ministral 3 Vision Language model | +| **Nemotron Nano V2 VL** | [nemotron-nano-v2-vl.md](/nemotron-nano-v2-vl) | NVIDIA Nemotron Nano V2 Vision Language model | +| **Qwen2.5 VL** | [qwen2.5-vl.md](/qwen2.5-vl) | Alibaba Cloud Qwen2.5 Vision Language model | +| **Qwen3 VL** | [qwen3-vl.md](/qwen3-vl) | Alibaba Cloud Qwen3 Vision Language model | + +## Quick Navigation + +### I want to + +**🔍 Find a specific VLM model** +→ Browse the model list above or use the [index page](/index) + +**🔄 Convert models between formats** +→ Each model page includes conversion examples for Hugging Face ↔ Megatron Bridge + +**🚀 Get started with training** +→ See [Training Documentation](/../../training/README) for training guides + +**📚 Understand VLM architecture** +→ Each model page documents vision-language architecture features + +**🔧 Add support for a new VLM** +→ Refer to [Adding New Models](/../../adding-new-models) + +## Related Documentation + +- **[Models Overview](/../README)** - Return to main models documentation +- **[Large Language Models](/../llm/README)** - LLM model documentation +- **[Training Documentation](/../../training/README)** - Training and customization guides +- **[Bridge Guide](/../../bridge-guide)** - Working with Hugging Face models +- **[Adding New Models](/../../adding-new-models)** - Extending model support + +## Vision Language Model Features + +VLMs typically support: + +- **Image Understanding** - Processing and understanding visual inputs +- **Multimodal Fusion** - Combining vision and language representations +- **Vision-Language Tasks** - Image captioning, visual question answering, and more +- **Cross-Modal Learning** - Learning relationships between visual and textual data + +--- + +**Ready to explore?** Choose a model from the list above or return to the [main documentation](/../../README). diff --git a/fern/v0.2.0/pages/models/vlm/gemma3-vl.mdx b/fern/v0.2.0/pages/models/vlm/gemma3-vl.mdx new file mode 100644 index 0000000000..78c5eaafe3 --- /dev/null +++ b/fern/v0.2.0/pages/models/vlm/gemma3-vl.mdx @@ -0,0 +1,64 @@ +--- +title: "Gemma 3 VL (Vision-Language)" +description: "" +--- + +[Google's Gemma 3 VL](https://huggingface.co/collections/google/gemma-3-release) is a family of vision-language models built on the same research and technology used to create Gemini models. The Gemma 3 VL architecture combines the text-generation capabilities of Gemma 3 with a SigLIP vision encoder for robust visual understanding. + +Gemma 3 VL models support multimodal tasks including image captioning, visual question answering, OCR, and general vision-language understanding. + +Gemma family models are supported via the Bridge system with auto-detected configuration and weight mapping. + +## Available Models + +### Vision-Language Models +- **Gemma 3 VL 4B** (`google/gemma-3-4b-it`): 4B parameter vision-language model + - 34 layers, 2560 hidden size + - 16 attention heads, 4 query groups (GQA) + - Vision encoder: SigLIP with 729M parameters + - Recommended: 1 node, 8 GPUs + +- **Gemma 3 VL 12B** (`google/gemma-3-12b-it`): 12B parameter vision-language model + - 48 layers, 3840 hidden size + - 24 attention heads, 8 query groups (GQA) + - Vision encoder: SigLIP with 729M parameters + - Recommended: 1 node, 8 GPUs + +- **Gemma 3 VL 27B** (`google/gemma-3-27b-it`): 27B parameter vision-language model + - 62 layers, 5376 hidden size + - 32 attention heads, 16 query groups (GQA) + - Vision encoder: SigLIP with 729M parameters + - Recommended: 2 nodes, 16 GPUs + +All models support a sequence length of 131,072 tokens and use hybrid attention patterns (sliding window + global). + +## Model Architecture Features + +Gemma 3 VL builds on the Gemma 3 architecture with additional multimodal capabilities: + +**Language Model Features:** +- **Hybrid Attention Pattern**: Alternates between global and local sliding window attention for efficient long-context processing +- **GeGLU Activation**: Uses gated linear units with GELU activation for improved performance +- **RMSNorm**: Layer normalization without mean centering for faster computation +- **Rotary Embeddings**: Separate RoPE configurations for local and global attention layers + +**Vision-Language Features:** +- **SigLIP Vision Encoder**: Pre-trained vision encoder with 729M parameters for robust visual understanding +- **Multimodal Integration**: Seamless integration of visual and textual information through learned projection layers +- **Flexible Image Handling**: Supports variable resolution images and multiple images per conversation + +## Examples + +For checkpoint conversion, inference, finetuning recipes, and step-by-step training guides, see the [Gemma 3 VL Examples](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/models/vlm/gemma3_vl/README.md). + +## Hugging Face Model Cards + +- Gemma 3 VL 4B: https://huggingface.co/google/gemma-3-4b-it +- Gemma 3 VL 12B: https://huggingface.co/google/gemma-3-12b-it +- Gemma 3 VL 27B: https://huggingface.co/google/gemma-3-27b-it + +## Related Docs +- Text-Only Models: [Gemma 3](/../llm/gemma3) +- Recipe usage: [Recipe usage](/../../recipe-usage) +- Customizing the training recipe configuration: [Configuration overview](/../../training/config-container-overview) +- Training entry points: [Entry points](/../../training/entry-points) diff --git a/fern/v0.2.0/pages/models/vlm/glm-45v.mdx b/fern/v0.2.0/pages/models/vlm/glm-45v.mdx new file mode 100644 index 0000000000..82c777a919 --- /dev/null +++ b/fern/v0.2.0/pages/models/vlm/glm-45v.mdx @@ -0,0 +1,57 @@ +--- +title: "GLM-4.5V" +description: "" +--- + +[GLM-4.5V](https://huggingface.co/zai-org/GLM-4.5V) is a powerful vision-language model built on the GLM-4.5 Air architecture. GLM-4.5V combines a 106B parameter sparse MoE language model with a vision encoder for robust multimodal understanding of images and videos. + +GLM-4.5V supports multimodal tasks including image captioning, visual question answering, OCR, video understanding, and general vision-language reasoning. The model leverages Multi-Resolution Rotary Position Embedding (MRoPE) for enhanced spatial understanding. + +GLM family models are supported via the Bridge system with auto-detected configuration and weight mapping. + + +Please update `transformers` version to 4.57.1 or higher in order to use the GLM-4.5V model. + + +## Available Models + +### Vision-Language Models +- **GLM-4.5V** (`zai-org/GLM-4.5V`): 106B parameter vision-language model (based on GLM-4.5 Air) + - 46 decoder layers, 4096 hidden size + - 96 attention heads, 8 query groups (GQA) + - 128 MoE experts with shared experts + - ~12B active parameters per token + - Sequence length: 131,072 tokens + - Recommended: 32 nodes, 256 GPUs (LoRA/DoRA) or 64 nodes, 512 GPUs (Full SFT) + +## Model Architecture Features + +GLM-4.5V combines efficient sparse MoE language modeling with multimodal capabilities: + +**Language Model Features:** +- **Sparse MoE Architecture**: 128 routed experts with shared experts for efficient parameter usage +- **Grouped Query Attention (GQA)**: Memory-efficient attention with 8 query groups +- **SiLU Gated Linear Unit**: Gated linear units with SiLU activation for improved performance +- **RMSNorm**: Layer normalization without mean centering for faster computation +- **Multi-Resolution RoPE (MRoPE)**: Enhanced position embeddings with sections [8, 12, 12] for improved spatial understanding +- **Extended Context**: Supports up to 131,072 tokens + +**Vision-Language Features:** +- **Vision Encoder**: Pre-trained vision encoder for robust visual understanding +- **Multimodal Integration**: Seamless integration of visual and textual information +- **Image and Video Support**: Handles both static images and video inputs +- **Flexible Image Handling**: Supports variable resolution images and multiple images per conversation + +## Examples + +For checkpoint conversion, inference, finetuning recipes, and step-by-step training guides, see the [GLM-4.5V Examples](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/models/vlm/glm_45v/README.md). + +## Hugging Face Model Cards + +- GLM-4.5V: https://huggingface.co/zai-org/GLM-4.5V + +## Related Docs +- Related LLM: [GLM 4.5](/../llm/glm45) +- Recipe usage: [Recipe usage](/../../recipe-usage) +- Customizing the training recipe configuration: [Configuration overview](/../../training/config-container-overview) +- Training entry points: [Entry points](/../../training/entry-points) diff --git a/fern/v0.2.0/pages/models/vlm/index.mdx b/fern/v0.2.0/pages/models/vlm/index.mdx new file mode 100644 index 0000000000..3c78d0224e --- /dev/null +++ b/fern/v0.2.0/pages/models/vlm/index.mdx @@ -0,0 +1,6 @@ +--- +title: "Vision Language Models" +description: "" +--- + +This section documents Vision Language Models supported by Megatron Bridge. diff --git a/fern/v0.2.0/pages/models/vlm/ministral3.mdx b/fern/v0.2.0/pages/models/vlm/ministral3.mdx new file mode 100644 index 0000000000..60f224cc94 --- /dev/null +++ b/fern/v0.2.0/pages/models/vlm/ministral3.mdx @@ -0,0 +1,72 @@ +--- +title: "Ministral 3" +description: "" +--- + +[Mistral AI's Ministral 3](https://huggingface.co/collections/mistralai/ministral-3) is a family of edge-optimized vision-language models designed for deployment across various hardware configurations. The Ministral 3 architecture combines a powerful language model with a vision encoder for multimodal understanding. + +Ministral 3 models support multimodal tasks including image captioning, visual question answering, OCR, and general vision-language understanding. Despite their compact size, these models deliver strong performance for on-device and edge deployment scenarios. + +Ministral family models are supported via the Bridge system with auto-detected configuration and weight mapping. + + +Please upgrade to `transformers` v5 and upgrade `mistral-common` in order to use the Ministral 3 models. + + +## Available Models + +### Vision-Language Models +- **Ministral 3 3B** (`mistralai/Ministral-3-3B-Base-2512`): 3.4B parameter vision-language model + - 26 layers, 3072 hidden size + - 32 attention heads, 8 query groups (GQA) + - Vision encoder: ~0.4B parameters + - Recommended: 1 node, 8 GPUs + +- **Ministral 3 8B** (`mistralai/Ministral-3-8B-Base-2512`): 8.4B parameter vision-language model + - 34 layers, 4096 hidden size + - 32 attention heads, 8 query groups (GQA) + - Vision encoder: ~0.4B parameters + - Recommended: 1 node, 8 GPUs + +- **Ministral 3 14B** (`mistralai/Ministral-3-14B-Base-2512`): ~14B parameter vision-language model + - 40 layers, 5120 hidden size + - 32 attention heads, 8 query groups (GQA) + - Vision encoder: ~0.4B parameters + - Recommended: 1 node, 8 GPUs + +All models support extended context lengths up to 256K tokens using YaRN RoPE scaling. + +## Model Architecture Features + +Ministral 3 combines efficient language modeling with multimodal capabilities: + +**Language Model Features:** +- **YaRN RoPE Scaling**: Advanced rope scaling for extended context lengths (up to 256K tokens) +- **Grouped Query Attention (GQA)**: Memory-efficient attention mechanism with 8 query groups +- **SwiGLU Activation**: Gated linear units with SiLU activation for improved performance +- **RMSNorm**: Layer normalization without mean centering for faster computation +- **Llama 4 Attention Scaling**: Position-dependent attention scaling for improved long-context handling + +**Vision-Language Features:** +- **Vision Encoder**: Pre-trained vision encoder for robust visual understanding +- **Multimodal Projector**: Projects vision features to language model space +- **Flexible Image Handling**: Supports variable resolution images and multiple images per conversation + +## Examples + +For checkpoint conversion, inference, finetuning recipes, and step-by-step training guides, see the [Ministral 3 Examples](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/models/vlm/ministral3/README.md). + +## Hugging Face Model Cards + +- Ministral 3 3B Base: https://huggingface.co/mistralai/Ministral-3-3B-Base-2512 +- Ministral 3 3B Instruct: https://huggingface.co/mistralai/Ministral-3-3B-Instruct-2512 +- Ministral 3 8B Base: https://huggingface.co/mistralai/Ministral-3-8B-Base-2512 +- Ministral 3 8B Instruct: https://huggingface.co/mistralai/Ministral-3-8B-Instruct-2512 +- Ministral 3 14B Base: https://huggingface.co/mistralai/Ministral-3-14B-Base-2512 +- Ministral 3 14B Instruct: https://huggingface.co/mistralai/Ministral-3-14B-Instruct-2512 + +## Related Docs +- Related LLM: [Mistral](/../llm/mistral) +- Recipe usage: [Recipe usage](/../../recipe-usage) +- Customizing the training recipe configuration: [Configuration overview](/../../training/config-container-overview) +- Training entry points: [Entry points](/../../training/entry-points) diff --git a/fern/v0.2.0/pages/models/vlm/nemotron-nano-v2-vl.mdx b/fern/v0.2.0/pages/models/vlm/nemotron-nano-v2-vl.mdx new file mode 100644 index 0000000000..c6cc61b8d0 --- /dev/null +++ b/fern/v0.2.0/pages/models/vlm/nemotron-nano-v2-vl.mdx @@ -0,0 +1,197 @@ +--- +title: "Nemotron Nano V2 VL" +description: "" +--- + +NVIDIA Nemotron Nano v2 VL is an open 12B multimodal reasoning model for document intelligence and video understanding. +It enables [AI assistants](https://www.nvidia.com/en-us/use-cases/ai-assistants) to extract, interpret, and act on +information across text, images, tables, and videos. This makes the model valuable for agents focused on data analysis, +document processing and visual understanding in applications like generating reports, curating videos, and dense +captioning for media asset management, and retrieval-augmented search. + +NeMo Megatron Bridge supports finetuning this model (including LoRA finetuning) on single-image, multi-image, and video +datasets. +The finetuned model can be converted back to the 🤗 Hugging Face format for downstream evaluation. + + +Please use the custom container `nvcr.io/nvidia/nemo:25.09.nemotron_nano_v2_vl` when working with this model. + +Run all commands from `/opt/Megatron-Bridge` (e.g. `docker run -w /opt/Megatron-Bridge ...`) + + + +We use the following environment variables throughout this page +- `HF_MODEL_PATH=nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16` +- `MEGATRON_MODEL_PATH=/models/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16` (feel free to set your own path) + +Unless explicitly stated, any megatron model path in the commands below should NOT contain the iteration number +`iter_xxxxxx`. For more details on checkpointing, please see +[here](https://docs.nvidia.com/nemo/megatron-bridge/latest/training/checkpointing.html#checkpoint-contents) + + +## Conversion with 🤗 Hugging Face + +### Import HF → Megatron +To import the HF model to your desired `$MEGATRON_MODEL_PATH`, run the following command. +```bash +python examples/conversion/convert_checkpoints.py import \ +--hf-model $HF_MODEL_PATH \ +--megatron-path $MEGATRON_MODEL_PATH \ +--trust-remote-code +``` + +### Export Megatron → HF +You can export a trained model with the following command. +```bash +python examples/conversion/convert_checkpoints.py export \ +--hf-model $HF_MODEL_PATH \ +--megatron-path \ +--hf-path \ +--not-strict +``` + +Note: it is normal to see a warning that `vision_model.radio_model.input_conditioner.norm_mean` and `vision_model.radio_model.input_conditioner.norm_std` from source are not in the exported checkpoint. These two weights are not needed in the checkpoint. + +### Run In-Framework Inference on Converted Checkpoint +You can run a quick sanity check on the converted checkpoint with the following command. +```bash +python examples/conversion/hf_to_megatron_generate_vlm.py \ +--hf_model_path $HF_MODEL_PATH \ +--megatron_model_path $MEGATRON_MODEL_PATH \ +--image_path \ +--prompt "Describe this image." \ +--max_new_tokens 100 \ +--use_llava_model +``` + +Note: +- `--megatron_model_path` is optional. If not specified, the script will convert the model and then run forward. If + specified, the script will just load the megatron model +- `--max_new_tokens` controls the number of tokens to generate. +- For inference with multiple images, pass in a comma-separated list, e.g. + `--image_path="/path/to/example1.jpeg,/path/to/example2.jpeg"`. + Use a suitable prompt, e.g. `--prompt="Describe the two images in detail."`. +- For inference with video, pass in video path instead, e.g. `--video_path="/path/to/demo.mp4"`. Use a suitable prompt, + e.g. `--prompt="Describe what you see."`. + +## Finetuning Recipes +Before training, ensure the following environment variables are set. +1. `SAVE_DIR`: to specify a checkpoint and log saving directory, used in the commands below. +2. `HF_TOKEN`: to download models from HF Hub. +3. `HF_HOME`: (optional) to avoid re-downloading models and datasets every time. +4. `WANDB_API_KEY`: (optional) to enable WandB logging. + +### Full Finetuning +Example usage for full parameter finetuning using the +[Raven dataset](https://huggingface.co/datasets/HuggingFaceM4/the_cauldron/viewer/raven): + +```bash +torchrun --nproc-per-node=8 examples/models/vlm/nemotron_vl/finetune_nemotron_nano_v2_vl.py \ +--hf-model-path $HF_MODEL_PATH \ +--pretrained-checkpoint \ +dataset.maker_name=make_raven_dataset \ +logger.wandb_project= \ +logger.wandb_save_dir=$SAVE_DIR \ +checkpoint.save=$SAVE_DIR/ +``` + +Note: +- The config file `examples/models/vlm/nemotron_vl/conf/nemotron_nano_v2_vl_override_example.yaml` contains a list of arguments + that can be overridden in the command. For example, you can set `train.global_batch_size=` in the command. +- To change the dataset, you only need to change `dataset.maker_name`. See the dataset section below for details. +- After training, you can run inference with `hf_to_megatron_generate_vlm.py` by supplying the trained megatron checkpoint. + You can also export the trained checkpoint to Hugging Face format. +- This full finetuning recipe requires at least 4xH100 (80G) GPUs. + +### Parameter-Efficient Finetuning (PEFT) +Parameter-efficient finetuning (PEFT) using LoRA is supported. +LoRA can be independently applied to the vision model, vision projection, and language model. We support two commonly used +settings out of the box in the example script: +1. Apply LoRA to the language model, and fully finetune the vision model and projection (used when the visual + distribution is substantially different from pretrained.) + +```bash +torchrun --nproc-per-node=8 examples/models/vlm/nemotron_vl/finetune_nemotron_nano_v2_vl.py \ +--hf-model-path $HF_MODEL_PATH \ +--pretrained-checkpoint $MEGATRON_MODEL_PATH \ +--lora-on-language-model \ +dataset.maker_name=make_raven_dataset \ +logger.wandb_project= \ +logger.wandb_save_dir=$SAVE_DIR \ +checkpoint.save=$SAVE_DIR/ \ +model.freeze_language_model=True \ +model.freeze_vision_model=False \ +model.freeze_vision_projection=False +``` + +2. Apply LoRA to all linear layers in attention and MLP modules of the vision model, vision projection, and the language model. + +```bash +torchrun --nproc-per-node=8 examples/models/vlm/nemotron_vl/finetune_nemotron_nano_v2_vl.py \ +--hf-model-path $HF_MODEL_PATH \ +--pretrained-checkpoint $MEGATRON_MODEL_PATH \ +--lora-on-language-model \ +—-lora-on-vision-model \ +dataset.maker_name=make_raven_dataset \ +logger.wandb_project= \ +logger.wandb_save_dir=$SAVE_DIR \ +checkpoint.save=$SAVE_DIR/ \ +model.freeze_language_model=True \ +model.freeze_vision_model=True \ +model.freeze_vision_projection=True +``` + +These LoRA finetuning recipe requires at least 2xH100 (80G) GPUs. + +A LoRA checkpoint only contains the learnable adapter weights. In order to convert the LoRA checkpoint to Hugging Face +format for downstream evaluation, it is necessary to merge the LoRA adapters back to the base model. + +```bash +python examples/peft/merge_lora.py \ +--hf-model-path $HF_MODEL_PATH \ +--lora-checkpoint /iter_N \ +--output +``` +You can now run in-framework inference with `hf_to_megatron_generate_vlm.py` by supplying the merged LoRA checkpoint. +You can also export the merged LoRA checkpoint to Hugging Face format. + +## Example Datasets + +Megatron Bridge supports various vision-language dataset examples which can be used to finetune Nemotron Nano V2 VL: +| Dataset | Maker Name | Description | +|---------|------------|-------------| +| [cord-v2](https://huggingface.co/datasets/naver-clova-ix/cord-v2) | `make_cord_v2_dataset` | OCR receipts: Single-image-text dataset for receipt understanding, outputs xml-like annotated text. | +| [MedPix-VQA](https://huggingface.co/datasets/mmoukouba/MedPix-VQA) | `make_medpix_dataset` | Medical VQA: Single-image question-answer dataset covering clinical medical images and free-form answers. | +| [The Cauldron (Raven subset)](https://huggingface.co/datasets/HuggingFaceM4/the_cauldron) | `make_raven_dataset` | Visual reasoning: Multi-image, vision reasoning dataset for analogical reasoning in different visual layouts. | +| [LLaVA-Video-178K (0_30_s_nextqa subset)](https://huggingface.co/datasets/lmms-lab/LLaVA-Video-178K) | `make_llava_video_178k_dataset` | Video understanding: video question-answer dataset covering every-day scenarios. | + +`cord-v2` is a tiny dataset and is included for demonstration only. This dataset is not recommended for PEFT tuning for this model since the XML output format interacts with the special tokens, leading to unexpected results. + +Note on video training example: +- We provide a video config yaml file instead of the default config yaml file that overwrites a few commands. Please + pass in `--config-file "examples/models/vlm/nemotron_vl/conf/nemotron_nano_v2_vl_video.yaml"`. +- The LLaVA video dataset requires manual download beforehand. Please place the downloaded and extracted video files + in a folder `VIDEO_ROOT` and pass it in to the maker with `dataset.maker_kwargs={"video_root_path":$VIDEO_ROOT}`. + In the nextqa subset example, `VIDEO_ROOT` should look like + ``` + $VIDEO_ROOT/ + ├── NextQA/ + │ └── NExTVideo/ + │ └── 0000/ + │ └── 2440175990.mp4 + │ └── 0001/ + │ └── ... + └── ... + ``` + +Full video training example command: +```bash +torchrun --nproc-per-node=8 examples/models/vlm/nemotron_vl/finetune_nemotron_nano_v2_vl.py \ +--hf-model-path $HF_MODEL_PATH \ +--pretrained-checkpoint $MEGATRON_MODEL_PATH \ +--config-file "examples/models/vlm/nemotron_vl/conf/nemotron_nano_v2_vl_video.yaml" \ +logger.wandb_project= \ +logger.wandb_save_dir=$SAVE_DIR \ +checkpoint.save=$SAVE_DIR/ \ +dataset.maker_kwargs={"video_root_path":$VIDEO_ROOT} +``` diff --git a/fern/v0.2.0/pages/models/vlm/qwen2.5-vl.mdx b/fern/v0.2.0/pages/models/vlm/qwen2.5-vl.mdx new file mode 100644 index 0000000000..9e578070f6 --- /dev/null +++ b/fern/v0.2.0/pages/models/vlm/qwen2.5-vl.mdx @@ -0,0 +1,142 @@ +--- +title: "Qwen2.5-VL" +description: "" +--- + +Qwen2.5-VL is a series of vision-language models developed by Alibaba Cloud that enable multimodal understanding across text, images, and videos. The models support various vision-language tasks including image understanding, visual question answering, and multimodal reasoning. + +NeMo Megatron Bridge supports finetuning Qwen2.5-VL models (3B, 7B, 32B, and 72B variants) on single-image and multi-image datasets. +The finetuned model can be converted back to the 🤗 Hugging Face format for downstream evaluation. + + +We use the following environment variables throughout this page +- `HF_MODEL_PATH=Qwen/Qwen2.5-VL-3B-Instruct` (it can also be set to `Qwen/Qwen2.5-VL-7B-Instruct`, `Qwen/Qwen2.5-VL-32B-Instruct`, `Qwen/Qwen2.5-VL-72B-Instruct`) +- `MEGATRON_MODEL_PATH=/models/Qwen2.5-VL-3B-Instruct` (feel free to set your own path) + +Unless explicitly stated, any megatron model path in the commands below should NOT contain the iteration number +`iter_xxxxxx`. For more details on checkpointing, please see +[here](https://docs.nvidia.com/nemo/megatron-bridge/latest/training/checkpointing.html#checkpoint-contents) + + +## Conversion with 🤗 Hugging Face + +### Import HF → Megatron +To import the HF model to your desired `$MEGATRON_MODEL_PATH`, run the following command. +```bash +python examples/conversion/convert_checkpoints.py import \ +--hf-model $HF_MODEL_PATH \ +--megatron-path $MEGATRON_MODEL_PATH +``` + +### Export Megatron → HF +You can export a trained model with the following command. +```bash +python examples/conversion/convert_checkpoints.py export \ +--hf-model $HF_MODEL_PATH \ +--megatron-path \ +--hf-path +``` + +### Run In-Framework Inference on Converted Checkpoint +You can run a quick sanity check on the converted checkpoint with the following command. +```bash +python examples/conversion/hf_to_megatron_generate_vlm.py \ +--hf_model_path $HF_MODEL_PATH \ +--megatron_model_path $MEGATRON_MODEL_PATH \ +--image_path \ +--prompt "Describe this image." \ +--max_new_tokens 100 +``` + +Note: +- `--megatron_model_path` is optional. If not specified, the script will convert the model and then run forward. If + specified, the script will just load the megatron model +- `--max_new_tokens` controls the number of tokens to generate. +- You can also use image URLs: `--image_path="https://example.com/image.jpg"` + +## Finetuning Recipes +Before training, ensure the following environment variables are set. +1. `SAVE_DIR`: to specify a checkpoint and log saving directory, used in the commands below. +2. `HF_TOKEN`: to download models from HF Hub (if required). +3. `HF_HOME`: (optional) to avoid re-downloading models and datasets every time. +4. `WANDB_API_KEY`: (optional) to enable WandB logging. + +### Full Finetuning + +Example usage for full parameter finetuning: + +```bash +torchrun --nproc-per-node=8 examples/models/vlm/qwen_vl/finetune_qwen25_vl.py \ +--pretrained-checkpoint $MEGATRON_MODEL_PATH \ +--recipe qwen25_vl_3b_finetune_config \ +--dataset-type hf \ +dataset.maker_name=make_cord_v2_dataset \ +train.global_batch_size= \ +train.train_iters= \ +logger.wandb_project= \ +logger.wandb_save_dir=$SAVE_DIR \ +checkpoint.save=$SAVE_DIR/ +``` + +Note: +- The `--recipe` parameter selects the model size configuration. Available options: + - `qwen25_vl_3b_finetune_config` - for 3B model + - `qwen25_vl_7b_finetune_config` - for 7B model + - `qwen25_vl_32b_finetune_config` - for 32B model + - `qwen25_vl_72b_finetune_config` - for 72B model +- The config file `examples/models/vlm/qwen_vl/conf/qwen25_vl_pretrain_override_example.yaml` contains a list of arguments + that can be overridden in the command. For example, you can set `train.global_batch_size=` in the command. +- The dataset format should be JSONL with conversation format (see dataset section below). +- After training, you can run inference with `hf_to_megatron_generate_vlm.py` by supplying the trained megatron checkpoint. + You can also export the trained checkpoint to Hugging Face format. + +### Parameter-Efficient Finetuning (PEFT) +Parameter-efficient finetuning (PEFT) using LoRA or DoRA is supported. You can use the `--peft_scheme` argument to enable PEFT training: + +```bash +torchrun --nproc-per-node=8 examples/models/vlm/qwen_vl/finetune_qwen25_vl.py \ +--pretrained-checkpoint $MEGATRON_MODEL_PATH \ +--recipe qwen25_vl_3b_finetune_config \ +--peft_scheme lora \ +--dataset-type hf \ +dataset.maker_name=make_cord_v2_dataset \ +train.global_batch_size= \ +checkpoint.save=$SAVE_DIR/ +``` + +PEFT options: +- `--peft_scheme`: Set to `lora` for LoRA (Low-Rank Adaptation) or `dora` for DoRA (Weight-Decomposed Low-Rank Adaptation). Set to `None` or omit for full finetuning. + +You can also combine PEFT with freeze options to control which components are trainable: +- `model.freeze_language_model`: Set to `True` to freeze the language model +- `model.freeze_vision_model`: Set to `True` to freeze the vision encoder +- `model.freeze_vision_projection`: Set to `True` to freeze the vision projection layer + +Example with LoRA and freeze options: +```bash +torchrun --nproc-per-node=8 examples/models/vlm/qwen_vl/finetune_qwen25_vl.py \ +--pretrained-checkpoint $MEGATRON_MODEL_PATH \ +--recipe qwen25_vl_3b_finetune_config \ +--peft_scheme lora \ +model.freeze_language_model=True \ +model.freeze_vision_model=False \ +model.freeze_vision_projection=False \ +checkpoint.save=$SAVE_DIR/ +``` + +## Example Datasets + +Megatron Bridge supports various vision-language dataset examples which can be used to finetune Qwen 2.5 VL: +| Dataset | Maker Name | Description | +|---------|------------|-------------| +| [cord-v2](https://huggingface.co/datasets/naver-clova-ix/cord-v2) | `make_cord_v2_dataset` | OCR receipts: Single-image-text dataset for receipt understanding, outputs xml-like annotated text. | +| [MedPix-VQA](https://huggingface.co/datasets/mmoukouba/MedPix-VQA) | `make_medpix_dataset` | Medical VQA: Single-image question-answer dataset covering clinical medical images and free-form answers. | +| [The Cauldron (Raven subset)](https://huggingface.co/datasets/HuggingFaceM4/the_cauldron) | `make_raven_dataset` | Visual reasoning: Multi-image, vision reasoning dataset for analogical reasoning in different visual layouts. | + +To change the dataset, specify `dataset.maker_name=make_raven_dataset` + +## Hugging Face Model Cards +- Qwen2.5-VL-3B: `https://huggingface.co/Qwen/Qwen2.5-VL-3B-Instruct` +- Qwen2.5-VL-7B: `https://huggingface.co/Qwen/Qwen2.5-VL-7B-Instruct` +- Qwen2.5-VL-32B: `https://huggingface.co/Qwen/Qwen2.5-VL-32B-Instruct` +- Qwen2.5-VL-72B: `https://huggingface.co/Qwen/Qwen2.5-VL-72B-Instruct` diff --git a/fern/v0.2.0/pages/models/vlm/qwen3-vl.mdx b/fern/v0.2.0/pages/models/vlm/qwen3-vl.mdx new file mode 100644 index 0000000000..8db3995684 --- /dev/null +++ b/fern/v0.2.0/pages/models/vlm/qwen3-vl.mdx @@ -0,0 +1,26 @@ +--- +title: "Qwen3-VL" +description: "" +--- + +Qwen3-VL is the latest generation of vision-language models from Alibaba Cloud, supporting multimodal understanding across text, images, and videos. Qwen3-VL includes both dense models and Mixture-of-Experts (MoE) variants for improved efficiency. + +NeMo Megatron Bridge supports finetuning Qwen3-VL models (8B dense and 30B MoE variants). + + +We use the following environment variables throughout this page +- `HF_MODEL_PATH=Qwen/Qwen3-VL-8B-Instruct` (or `Qwen/Qwen3-VL-30B-A3B-Instruct` for MoE) +- `MEGATRON_MODEL_PATH=/models/Qwen3-VL-8B-Instruct` (feel free to set your own path) +Unless explicitly stated, any megatron model path in the commands below should NOT contain the iteration number +`iter_xxxxxx`. For more details on checkpointing, please see +[here](https://docs.nvidia.com/nemo/megatron-bridge/latest/training/checkpointing.html#checkpoint-contents) + + +## Examples + +For checkpoint conversion, inference, finetuning recipes, and step-by-step training guides, see the [Qwen3-VL Examples](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/models/vlm/qwen3_vl/README.md). + +## Hugging Face Model Cards +- Qwen3-VL-8B: `https://huggingface.co/Qwen/Qwen3-VL-8B-Instruct` +- Qwen3-VL-30B-A3B (MoE): `https://huggingface.co/Qwen/Qwen3-VL-30B-A3B-Instruct` +- Qwen3-VL-235B-A22B (MoE): `https://huggingface.co/Qwen/Qwen3-VL-235B-A22B-Instruct` diff --git a/fern/v0.2.0/pages/nemo2-migration-guide.mdx b/fern/v0.2.0/pages/nemo2-migration-guide.mdx new file mode 100644 index 0000000000..3b3ae058fd --- /dev/null +++ b/fern/v0.2.0/pages/nemo2-migration-guide.mdx @@ -0,0 +1,1791 @@ +--- +title: "NeMo 2.0 to Megatron Bridge Migration Guide" +description: "" +--- + +This guide helps you migrate from NeMo 2.0 training and recipes to Megatron Bridge. Megatron Bridge retains the Pythonic, code-first API that NeMo 2.0 developed while simplifying configuration into a single [`bridge.training.config.ConfigContainer`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.ConfigContainer.html) with typed sub-configs. Model parallelism and performance features from Megatron Core remain first-class. + +## What Stays the Same + +- **Megatron Core Foundation**: Megatron Bridge uses the same Megatron Core engine under the hood. +- **Model Parallelism**: Same TP/PP/CP/EP concepts with identical distributed training semantics. +- **High-Performance Features**: Mixed Precision, communication overlap, and other performance features are supported natively. +- **Pythonic API Retained**: Megatron Bridge preserves NeMo 2.0's philosophy of "configuration as code." + +## Model Configuration Mapping + +Megatron Bridge offers model providers that directly map to NeMo 2.0 model configs. + +### Examples + +| NeMo 2.0 | Megatron Bridge | +|----------|-----------------| +| `llm.Llama3Config8B` | [`bridge.models.Llama3ModelProvider8B`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.models.Llama3ModelProvider8B.html) | +| `llm.Llama31Config70B` | [`bridge.models.Llama31ModelProvider70B`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.models.Llama31ModelProvider70B.html) | +| `llm.Qwen2Config7B` | [`bridge.models.Qwen2ModelProvider7B`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.models.Qwen2ModelProvider7B.html) | +| `llm.DeepseekV2Config` | [`bridge.models.DeepseekV2ModelProvider`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.models.DeepseekV2ModelProvider.html) | + +### Supported Model Families + +Megatron Bridge supports the following model families with preset providers: +- **Base Models**: `GPTModelProvider`, `T5ModelProvider`, `MambaModelProvider` +- **Llama**: Llama2, Llama3, Llama3.1, Llama3.2, CodeLlama, Llama4 +- **Qwen**: Qwen2, Qwen2.5, Qwen3, Qwen3MoE, Qwen2.5VL +- **DeepSeek**: DeepSeek, DeepSeekV2, DeepSeekV2Lite, DeepSeekV3, Moonlight +- **Nemotron**: Nemotron3, Nemotron4, NemotronH, NemotronNano +- **NVIDIA Mamba**: Mamba variants and hybrid models + +For a complete list of all model providers and their parameters, see [`bridge.models`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.models.html). + +{/* TODO: Create a dedicated model support table with tested HF checkpoint mappings */} + +--- + +## Quick Start: Migration Examples + +This section shows complete migration examples for common training scenarios. For detailed configuration mappings, see [Configuration Migration](#configuration-migration). For entry point API details, see [Entry Points](#entry-points-pretrain-and-finetune). + +### Pretraining Migration Example + +#### Before: NeMo 2.0 +```python +from nemo import lightning as nl +from nemo.collections import llm +import nemo_run as run +from megatron.core.distributed import DistributedDataParallelConfig + +# Model configuration +model = run.Config( + llm.LlamaModel, + config=run.Config(llm.Llama3Config8B), # Preset config with all defaults +) + +# Strategy with parallelism settings +strategy = run.Config( + nl.MegatronStrategy, + tensor_model_parallel_size=2, + pipeline_model_parallel_size=2, + context_parallel_size=1, + sequence_parallel=False, + ddp=run.Config( + DistributedDataParallelConfig, + grad_reduce_in_fp32=True, + ), +) + +# Trainer setup +trainer = run.Config( + nl.Trainer, + max_steps=1000, + val_check_interval=100, + limit_val_batches=50, + log_every_n_steps=10, + devices=8, + num_nodes=1, + strategy=strategy, + plugins=nl.MegatronMixedPrecision(precision="bf16-mixed"), +) + +# Data configuration +data = run.Config( + llm.PreTrainingDataModule, + paths="/path/to/data_text_document", + seq_length=8192, + micro_batch_size=1, + global_batch_size=512, +) + +# Optimizer configuration +optim = llm.distributed_fused_adam_with_cosine_annealing( + max_lr=3e-4, + min_lr=3e-5, + warmup_steps=100, +) + +# Execute training +llm.pretrain(model, data, trainer, optim=optim) +``` + +#### Now: Megatron Bridge +```python +# Megatron Bridge configuration pattern +from megatron.bridge.training.config import ( + CheckpointConfig, + ConfigContainer, + GPTDatasetConfig, + LoggerConfig, + TrainingConfig, +) +from megatron.bridge.models import Llama3ModelProvider8B # Direct equivalent to Llama3Config8B +from megatron.core.optimizer import OptimizerConfig +from megatron.bridge.training.config import SchedulerConfig +from megatron.bridge.training.pretrain import pretrain +# Use the provided GPT forward step +from megatron.bridge.training.gpt_step import forward_step + +def create_config(): + return ConfigContainer( + # Model with parallelism built-in - using preset 8B config + model=Llama3ModelProvider8B( + # Parallelism settings (moved from MegatronStrategy) + tensor_model_parallel_size=2, + pipeline_model_parallel_size=2, + context_parallel_size=1, + sequence_parallel=False, + # Can still override any model params if needed + seq_length=8192, + ), + # Training loop configuration + train=TrainingConfig( + global_batch_size=512, + micro_batch_size=1, + train_iters=1000, # was max_steps + eval_interval=100, # was val_check_interval + eval_iters=50, # was limit_val_batches + ), + # Optimization and scheduling + optimizer=OptimizerConfig( + optimizer="adam", + lr=3e-4, + min_lr=3e-5, + use_distributed_optimizer=True, + ), + scheduler=SchedulerConfig( + lr_decay_style="cosine", + lr_warmup_iters=100, + lr_decay_iters=1000, + ), + # Data configuration + dataset=GPTDatasetConfig( + blend=["/path/to/data_text_document"], + seq_length=8192, + ), + # Checkpointing and logging + checkpoint=CheckpointConfig( + save="/path/to/checkpoints", + save_interval=100, + ckpt_format="torch_dist", + ), + logger=LoggerConfig(log_interval=10), # was log_every_n_steps + # Mixed precision + mixed_precision="bf16_mixed", + ) + +# Execute training +cfg = create_config() + +pretrain(cfg, forward_step_func=forward_step) +``` + +### Fine-Tuning Migration Example (SFT/PEFT) + +For fine-tuning, use [`bridge.training.config.FinetuningDatasetConfig`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.FinetuningDatasetConfig.html) for data and set `checkpoint.pretrained_checkpoint` to the base model. Optionally add a `peft` configuration for parameter-efficient training. + +#### Before: NeMo 2.0 +```python +from nemo import lightning as nl +from nemo.collections import llm +import nemo_run as run + +# Model and trainer configuration +model = run.Config(llm.LlamaModel, config=run.Config(llm.Llama3Config8B)) +trainer = run.Config( + nl.Trainer, + max_steps=500, + val_check_interval=100, + devices=8, + num_nodes=1, +) + +# Data configuration +data = run.Config( + llm.FineTuningDataModule, + dataset_root="/path/to/sft/data", + seq_length=2048, + micro_batch_size=1, + global_batch_size=128, +) + +# PEFT configuration +lora = llm.peft.LoRA( + target_modules=['linear_qkv', 'linear_proj'], + dim=32, + alpha=16, +) + +# Execute fine-tuning with PEFT +llm.finetune( + model=model, + data=data, + trainer=trainer, + peft=lora, + tokenizer="model", +) +``` + +#### Now: Megatron Bridge +```python +# Megatron Bridge fine-tuning configuration (with optional PEFT) +from megatron.bridge.models import Llama3ModelProvider8B +from megatron.bridge.peft import LoRA + +def create_finetune_config(): + return ConfigContainer( + model=Llama3ModelProvider8B( + # Preset config matching Llama3Config8B + ), + train=TrainingConfig( + micro_batch_size=1, + global_batch_size=128, + train_iters=500, + ), + # Finetuning dataset instead of pretraining dataset + dataset=FinetuningDatasetConfig( + dataset_root="/path/to/sft/data", + seq_length=2048, + do_validation=True, + do_test=True, + # Optional: packed sequence support + packed_sequence_specs=PackedSequenceSpecs( + packed_sequence_size=2048, + ), + ), + # Must specify pretrained checkpoint + checkpoint=CheckpointConfig( + pretrained_checkpoint="/path/to/pretrained/model", + save="/path/to/sft/checkpoints", + load="/path/to/sft/checkpoints", + save_interval=50, + ), + # Optional: Enable PEFT + peft=LoRA( + target_modules=["linear_qkv", "linear_proj"], + dim=32, + alpha=16, + ), + # ... other configs + ) +``` + +--- + +## Recipe Migration + +NeMo 2.0 and Megatron Bridge both provide pre-built recipes for popular models. In NeMo 2.0, recipes return `run.Partial` configurations. Megatron Bridge recipes return `ConfigContainer` objects. + +### Using Pre-Built Recipes + +Both frameworks offer ready-to-use recipes that you can customize: + +**NeMo 2.0**: Recipes in `nemo.collections.llm.recipes/` +```python +from nemo.collections import llm + +# Use pre-built recipe +recipe = llm.llama3_8b.pretrain_recipe(name="my_run", num_nodes=2) +``` + +**Megatron Bridge**: Recipes in `megatron.bridge.recipes/` +```python +from megatron.bridge.recipes.llama.llama3_8b import pretrain_config +from megatron.bridge.training import pretrain +from megatron.bridge.training.gpt_step import forward_step + +# Use pre-built recipe +cfg = pretrain_config() + +# Customize as needed +cfg.train.train_iters = 10000 +cfg.model.tensor_model_parallel_size = 4 + +# Launch training +pretrain(cfg, forward_step_func=forward_step) +``` + +For details on using and customizing recipes, see [Recipe Usage](/recipe-usage). + +### Migrating a Custom Recipe + +If you've created a custom NeMo 2.0 recipe, here's how to migrate it to Megatron Bridge: + +#### Before: NeMo 2.0 Recipe Structure + +```python +# nemo/collections/llm/recipes/llama3_8b.py +import nemo_run as run +from nemo import lightning as nl +from nemo.collections import llm + +@run.cli.factory(name="llama3_8b") +def model() -> run.Config[pl.LightningModule]: + return run.Config(llm.LlamaModel, config=run.Config(llm.Llama3Config8B)) + +def trainer( + tensor_parallelism: int = 1, + pipeline_parallelism: int = 1, + num_nodes: int = 1, + num_gpus_per_node: int = 8, + max_steps: int = 1000, +) -> run.Config[nl.Trainer]: + strategy = run.Config( + nl.MegatronStrategy, + tensor_model_parallel_size=tensor_parallelism, + pipeline_model_parallel_size=pipeline_parallelism, + ) + return run.Config( + nl.Trainer, + devices=num_gpus_per_node, + num_nodes=num_nodes, + max_steps=max_steps, + strategy=strategy, + val_check_interval=100, + limit_val_batches=50, + ) + +@run.cli.factory(target=llm.pretrain, name="llama3_8b") +def pretrain_recipe( + dir: Optional[str] = None, + name: str = "default", + num_nodes: int = 1, + num_gpus_per_node: int = 8, +) -> run.Partial: + return run.Partial( + llm.pretrain, + model=model(), + trainer=trainer(num_nodes=num_nodes, num_gpus_per_node=num_gpus_per_node), + data=run.Config( + llm.PreTrainingDataModule, + paths="/path/to/data_text_document", + seq_length=8192, + global_batch_size=512, + micro_batch_size=1, + ), + log=llm.default_log(dir=dir, name=name), + optim=llm.distributed_fused_adam_with_cosine_annealing(max_lr=3e-4), + resume=llm.default_resume(), + ) + +# Usage +if __name__ == "__main__": + recipe = pretrain_recipe(name="my_run", num_nodes=2) + # Submitted via nemo-run or executed directly +``` + +#### Now: Megatron Bridge Recipe Structure + +```python +# my_recipes/llama3_8b.py +from typing import Optional +from megatron.bridge.training.config import ( + ConfigContainer, + TrainingConfig, + GPTDatasetConfig, + CheckpointConfig, + SchedulerConfig, +) +from megatron.core.optimizer import OptimizerConfig +from megatron.bridge.models import Llama3ModelProvider8B +from megatron.bridge.training import pretrain + +def llama3_8b_config( + # Model/parallelism params + tensor_parallelism: int = 1, + pipeline_parallelism: int = 1, + # Training params + train_iters: int = 1000, + eval_interval: int = 100, + eval_iters: int = 50, + # Data params + data_path: str = "/path/to/data_text_document", + seq_length: int = 8192, + global_batch_size: int = 512, + micro_batch_size: int = 1, + # Checkpoint params + checkpoint_dir: Optional[str] = None, + save_interval: int = 1000, +) -> ConfigContainer: + """Create a Llama3 8B pretraining configuration.""" + return ConfigContainer( + model=Llama3ModelProvider8B( + # Preset architecture from Llama3Config8B (num_layers=32, hidden_size=4096, etc.) + # Only need to specify parallelism and overrides + tensor_model_parallel_size=tensor_parallelism, + pipeline_model_parallel_size=pipeline_parallelism, + ), + train=TrainingConfig( + train_iters=train_iters, + eval_interval=eval_interval, + eval_iters=eval_iters, + global_batch_size=global_batch_size, + micro_batch_size=micro_batch_size, + ), + dataset=GPTDatasetConfig( + blend=[data_path], + seq_length=seq_length, + ), + optimizer=OptimizerConfig( + optimizer="adam", + lr=3e-4, + use_distributed_optimizer=True, + ), + scheduler=SchedulerConfig( + lr_decay_style="cosine", + lr_warmup_iters=100, + ), + checkpoint=CheckpointConfig( + save=checkpoint_dir or "/results/checkpoints", + save_interval=save_interval, + ), + mixed_precision="bf16-mixed", + ) + +# Usage +if __name__ == "__main__": + from megatron.bridge.training.gpt_step import forward_step + + cfg = llama3_8b_config( + train_iters=10000, + checkpoint_dir="/my/checkpoints", + tensor_parallelism=2, + ) + pretrain(cfg, forward_step_func=forward_step) +``` + +**Migration steps:** +1. Replace `run.Partial` with a function returning `ConfigContainer` +2. Move all `trainer`, `strategy`, and distributed settings into model provider +3. Consolidate `log`, `optim`, `resume` into respective config objects +4. Remove `@run.cli.factory` decorators (optional: use your own CLI framework) +5. Launch with `torchrun` or similar launcher—device count no longer passed to training function + +--- + +## Configuration Migration + +### Overview +What used to be configured across Lightning `Trainer` arguments, callbacks, and `MegatronStrategy` parameters is now centralized into a set of configuration classes: + +| Configuration Area | Megatron Bridge Config Class | +|-------------------|-------------------| +| Training loop settings | [`bridge.training.config.TrainingConfig`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.TrainingConfig.html) | +| Checkpointing | [`bridge.training.config.CheckpointConfig`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.CheckpointConfig.html) | +| Logging and monitoring | [`bridge.training.config.LoggerConfig`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.LoggerConfig.html) | +| Distributed training initialization | [`bridge.training.config.DistributedInitConfig`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.DistributedInitConfig.html) | +| Mixed precision | [`bridge.training.mixed_precision.MixedPrecisionConfig`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.mixed_precision.MixedPrecisionConfig.html) | +| Performance profiling | [`bridge.training.config.ProfilingConfig`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.ProfilingConfig.html) | + +For detailed documentation on each configuration area, see the training documentation: +- [Config Container Overview](/training/config-container-overview) - Overview of the configuration system +- [Training Loop Settings](/training/training-loop-settings) - Training loop parameters and validation +- [Checkpointing](/training/checkpointing) - Checkpointing and model persistence +- [Optimizer Scheduler](/training/optimizer-scheduler) - Optimization and learning rate scheduling +- [Logging](/training/logging) - Logging, TensorBoard, and Weights & Biases +- [Profiling](/training/profiling) - Performance profiling with Nsys and PyTorch + +### Training Configuration Migration +Lightning `Trainer` parameters are now managed through dedicated configuration classes. + +| **Setting Category** | **NeMo 2.0 Location** | **Megatron Bridge Location** | **Details** | +|---------------------|----------------------|-------------------|-------------| +| **Training iterations** | `trainer.max_steps` | [`bridge.training.config.TrainingConfig.train_iters`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.TrainingConfig.train_iters.html) | Total training iterations | +| **Validation frequency** | `trainer.val_check_interval` | [`bridge.training.config.TrainingConfig.eval_interval`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.TrainingConfig.eval_interval.html) | Steps between validation runs | +| **Validation iterations** | `trainer.limit_val_batches` | [`bridge.training.config.TrainingConfig.eval_iters`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.TrainingConfig.eval_iters.html) | Number of validation steps per run | +| **Test iterations** | `trainer.limit_test_batches` | [`bridge.training.config.TrainingConfig.eval_iters`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.TrainingConfig.eval_iters.html) | Number of test steps (shares eval_iters) | +| **Logging frequency** | `trainer.log_every_n_steps` | [`bridge.training.config.LoggerConfig.log_interval`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.LoggerConfig.log_interval.html) | Logging frequency | + +#### Before: NeMo 2.0 +```python +trainer = run.Config( + nl.Trainer, + max_steps=1000, + val_check_interval=100, # validation frequency + limit_val_batches=50, # validation iterations per run + limit_test_batches=100, # test iterations + log_every_n_steps=10, +) +``` + +#### Now: Megatron Bridge +```python +train_config = TrainingConfig( + train_iters=1000, # was max_steps + eval_interval=100, # was val_check_interval + eval_iters=50, # was limit_val_batches (for both val and test) +) +logger_config = LoggerConfig(log_interval=10) # was log_every_n_steps +``` + +### Data Configuration Migration + +NeMo 2.0 uses `PreTrainingDataModule` and `FineTuningDataModule` classes. Megatron Bridge uses configuration objects: [`bridge.training.config.GPTDatasetConfig`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.GPTDatasetConfig.html) for pretraining and [`bridge.training.config.FinetuningDatasetConfig`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.FinetuningDatasetConfig.html) for fine-tuning. + +#### Pretraining Data + +##### Before: NeMo 2.0 PreTrainingDataModule + +```python +from nemo.collections.llm.gpt.data import PreTrainingDataModule + +# Single dataset +data = PreTrainingDataModule( + paths="/path/to/train_data_text_document", + seq_length=4096, + micro_batch_size=1, + global_batch_size=512, + num_workers=8, + split="949,50,1", # train/val/test split ratios +) + +# Multiple datasets with weights +data = PreTrainingDataModule( + paths=["30", "/path/to/dataset1_text_document", + "70", "/path/to/dataset2_text_document"], + seq_length=4096, + micro_batch_size=1, + global_batch_size=512, + split="949,50,1", +) + +# Separate train/val/test datasets +data = PreTrainingDataModule( + paths={ + "train": ["/path/to/train_data_text_document"], + "validation": ["/path/to/val_data_text_document"], + "test": ["/path/to/test_data_text_document"], + }, + seq_length=4096, + micro_batch_size=1, + global_batch_size=512, +) +``` + +##### Now: Megatron Bridge GPTDatasetConfig + +```python +from megatron.bridge.training.config import GPTDatasetConfig, TrainingConfig + +# Single dataset +dataset_config = GPTDatasetConfig( + blend=["/path/to/train_data_text_document"], + seq_length=4096, + split="949,50,1", +) +train_config = TrainingConfig( + micro_batch_size=1, + global_batch_size=512, +) + +# Multiple datasets with weights (blending) +dataset_config = GPTDatasetConfig( + blend=[ + "/path/to/dataset1_text_document", + "/path/to/dataset2_text_document", + ], + blend_weights=[0.3, 0.7], # Explicit weights (not zipped with paths) + seq_length=4096, + split="949,50,1", +) +``` + +**Key differences:** +- NeMo 2.0's `paths` → Megatron Bridge's `blend` +- NeMo 2.0's zipped list `["30", "path1", "70", "path2"]` → Megatron Bridge's separate `blend` and `blend_weights` +- Batch sizes move from data module to `TrainingConfig` +- Dataloader options (`num_workers`, `pin_memory`, etc.) available in both configs + +#### Fine-Tuning Data + +##### Before: NeMo 2.0 FineTuningDataModule + +```python +from nemo.collections.llm.gpt.data import FineTuningDataModule + +data = FineTuningDataModule( + dataset_root="/path/to/instruction_data", + seq_length=2048, + micro_batch_size=1, + global_batch_size=128, + num_workers=8, +) +``` + +##### Now: Megatron Bridge FinetuningDatasetConfig + +```python +from megatron.bridge.training.config import FinetuningDatasetConfig, TrainingConfig + +dataset_config = FinetuningDatasetConfig( + dataset_root="/path/to/instruction_data", + seq_length=2048, + do_validation=True, + do_test=False, + # Dataloader options (inherited from DataloaderConfig) + num_workers=8, + pin_memory=True, + persistent_workers=False, +) +train_config = TrainingConfig( + micro_batch_size=1, + global_batch_size=128, +) +``` + +**Key differences:** +- Batch sizes move to `TrainingConfig` +- Explicit control over finetuning validation/test splits via `do_validation` and `do_test` +- Dataloader options (`num_workers`, `pin_memory`, etc.) available via `FinetuningDatasetConfig` + +### Tokenizer Migration + +Megatron Bridge uses [`bridge.training.tokenizers.config.TokenizerConfig`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.tokenizers.config.TokenizerConfig.html) for consistent tokenizer setup across different model types. + +#### Before: NeMo 2.0 +```python +# Option 1: Using get_nmt_tokenizer utility +from nemo.collections.nlp.modules.common.tokenizer_utils import get_nmt_tokenizer + +tokenizer = get_nmt_tokenizer( + library="megatron", + model_name="GPT2BPETokenizer", + vocab_file="/path/to/vocab.json", + merges_file="/path/to/merges.txt", +) + +# Option 2: Using run.Config with tokenizer classes +import nemo_run as run +from nemo.collections.common.tokenizers.huggingface.auto_tokenizer import AutoTokenizer + +tokenizer = run.Config( + AutoTokenizer, + pretrained_model_name="meta-llama/Llama-3-8B", +) +``` + +#### Now: Megatron Bridge +```python +# Dedicated tokenizer configuration +from megatron.bridge.training.tokenizers.config import TokenizerConfig + +# GPT2 BPE Tokenizer +tokenizer_config = TokenizerConfig( + tokenizer_type="GPT2BPETokenizer", + vocab_file="/path/to/vocab.json", + merge_file="/path/to/merges.txt", +) + +# HuggingFace Tokenizer +tokenizer_config = TokenizerConfig( + tokenizer_type="HuggingFaceTokenizer", + tokenizer_model="meta-llama/Llama-3-8B", +) +``` + +#### Vocab Size Priority + +In Megatron Bridge, vocabulary size can be specified in either the model provider or derived from the tokenizer. The priority order is: + +1. **Model provider `vocab_size` is set**: Uses the model's vocab size + - Must be `>= tokenizer.vocab_size` (raises error if smaller) + - Sets `should_pad_vocab=False` (no automatic padding) + - Useful when you need a specific vocab size (e.g., for checkpoint compatibility) + +2. **Model provider `vocab_size` is None**: Uses tokenizer's vocab size + - Automatically derived from `tokenizer.vocab_size` after building the tokenizer. + - Sets `should_pad_vocab=True` (enables padding for efficient parallelism) + +```python +# Option 1: Let tokenizer determine vocab size +config = ConfigContainer( + model=Llama3ModelProvider8B( + # vocab_size not set - will use tokenizer's vocab size + vocab_size=None, + ), + tokenizer=TokenizerConfig( + tokenizer_type="HuggingFaceTokenizer", + tokenizer_model="meta-llama/Llama-3-8B", + ), +) + +# Option 2: Explicitly set vocab size in model +config = ConfigContainer( + model=Llama3ModelProvider8B( + vocab_size=128256, # Explicitly set (must be >= tokenizer vocab size) + ), + tokenizer=TokenizerConfig(...), +) +``` + +### Parallelism Configuration Migration +In NeMo 2.0, parallelism settings were configured on `MegatronStrategy`. In Megatron Bridge, these are set directly on the model provider: + +| **Parallelism Type** | **NeMo 2.0** | **Megatron Bridge** | +|---------------------|-------------|-----------| +| **Tensor Parallel** | `strategy.tensor_model_parallel_size` | `model.tensor_model_parallel_size` | +| **Pipeline Parallel** | `strategy.pipeline_model_parallel_size` | `model.pipeline_model_parallel_size` | +| **Virtual Pipeline** | `strategy.virtual_pipeline_model_parallel_size` | `model.virtual_pipeline_model_parallel_size` | +| **Microbatch Group Size** | `strategy.microbatch_group_size_per_vp_stage` | `model.microbatch_group_size_per_vp_stage` | +| **Pipeline Layer Distribution** | `strategy.num_layers_in_first_pipeline_stage` | `model.num_layers_in_first_pipeline_stage` | +| **Pipeline Layer Distribution** | `strategy.num_layers_in_last_pipeline_stage` | `model.num_layers_in_last_pipeline_stage` | +| **Context Parallel** | `strategy.context_parallel_size` | `model.context_parallel_size` | +| **Sequence Parallel** | `strategy.sequence_parallel` | `model.sequence_parallel` | +| **Expert Parallel** | `strategy.expert_model_parallel_size` | `model.expert_model_parallel_size` | +| **Expert Tensor Parallel** | `strategy.expert_tensor_parallel_size` | `model.expert_tensor_parallel_size` | +| **Pipeline Layout** | `strategy.pipeline_model_parallel_layout` | `model.pipeline_model_parallel_layout` | +| **Pipeline Comm Backend** | `strategy.pipeline_model_parallel_comm_backend` | `model.pipeline_model_parallel_comm_backend` | +| **Pipeline Dtype** | `strategy.pipeline_dtype` | `model.pipeline_dtype` | +| **Encoder Tensor Parallel** | `strategy.encoder_tensor_model_parallel_size` | `model.encoder_tensor_model_parallel_size` | +| **Encoder Pipeline Parallel** | `strategy.encoder_pipeline_model_parallel_size` | `model.encoder_pipeline_model_parallel_size` | +| **Embedding in Pipeline** | `strategy.account_for_embedding_in_pipeline_split` | `model.account_for_embedding_in_pipeline_split` | +| **Loss in Pipeline** | `strategy.account_for_loss_in_pipeline_split` | `model.account_for_loss_in_pipeline_split` | +| **TE RNG Tracker** | `strategy.use_te_rng_tracker` | `model.use_te_rng_tracker` | + +#### Before: NeMo 2.0 +```python +strategy = run.Config( + MegatronStrategy, + tensor_model_parallel_size=8, + pipeline_model_parallel_size=2, + context_parallel_size=2, + sequence_parallel=True, +) +``` + +#### Now: Megatron Bridge +```python +model = GPTModelProvider( + # Model architecture + num_layers=32, + hidden_size=4096, + # Parallelism co-located with model + tensor_model_parallel_size=8, + pipeline_model_parallel_size=2, + context_parallel_size=2, + sequence_parallel=True, +) +``` + +### DDP Configuration Migration +Some `MegatronStrategy` parameters move to [`bridge.training.config.DistributedDataParallelConfig`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.DistributedDataParallelConfig.html): + +| **Setting** | **NeMo 2.0** | **Megatron Bridge** | +|-------------|-------------|-----------| +| **Distributed Optimizer Instances** | `strategy.num_distributed_optimizer_instances` | [`bridge.training.config.DistributedDataParallelConfig.num_distributed_optimizer_instances`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.DistributedDataParallelConfig.num_distributed_optimizer_instances.html) | + +### Strategy Settings Migration +Additional `MegatronStrategy` parameters move to [`bridge.training.config.DistributedInitConfig`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.DistributedInitConfig.html): + +| **Setting** | **NeMo 2.0** | **Megatron Bridge** | +|-------------|-------------|-----------| +| **Process Groups** | `strategy.use_gloo_process_groups` | [`bridge.training.config.DistributedInitConfig.use_gloo_process_groups`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.DistributedInitConfig.use_gloo_process_groups.html) | +| **SHARP** | `strategy.use_sharp` | [`bridge.training.config.DistributedInitConfig.use_sharp`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.DistributedInitConfig.use_sharp.html) | +| **NCCL Config** | `strategy.nccl_communicator_config_path` | [`bridge.training.config.DistributedInitConfig.nccl_communicator_config_path`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.DistributedInitConfig.nccl_communicator_config_path.html) | +| **Mapping Order** | `strategy.use_tp_pp_dp_mapping` | [`bridge.training.config.DistributedInitConfig.use_tp_pp_dp_mapping`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.DistributedInitConfig.use_tp_pp_dp_mapping.html) | +| **Lazy Init** | `strategy.lazy_init` | [`bridge.training.config.DistributedInitConfig.lazy_init`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.DistributedInitConfig.lazy_init.html) | + +### Mixed Precision Migration +Mixed precision in NeMo 2.0 is controlled via precision plugins passed to the trainer. In Megatron Bridge, this moves to a dedicated configuration class: + +#### Before: NeMo 2.0 +```python +# Mixed precision via plugin +from nemo.lightning.pytorch.plugins import MegatronMixedPrecisionPlugin + +trainer = run.Config( + nl.Trainer, + plugins=[MegatronMixedPrecisionPlugin(precision="bf16-mixed")] +) +``` + +#### Now: Megatron Bridge +```python +# Option 1: Use preset strings +config = ConfigContainer( + mixed_precision="bf16_mixed", # Simple preset + # ... other configs +) + +# Option 2: Detailed configuration +config = ConfigContainer( + mixed_precision=MixedPrecisionConfig( + fp16=False, + bf16=True, + params_dtype=torch.bfloat16, + pipeline_dtype=torch.bfloat16, + ), + # ... other configs +) +``` + + +The `mixed_precision` configuration automatically synchronizes precision settings across model, optimizer, and DDP configurations, overriding any conflicting settings. This ensures consistent precision behavior throughout training. For details on configuration precedence and available recipes, see [Mixed Precision](/training/mixed-precision). + + +### Checkpointing Configuration Migration +Checkpointing configuration moves from `MegatronStrategy` parameters and `ModelCheckpoint` callback to [`bridge.training.config.CheckpointConfig`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.CheckpointConfig.html): + +| **Checkpoint Setting** | **NeMo 2.0** | **Megatron Bridge** | +|------------------------|-------------|-----------| +| **Save directory** | `ModelCheckpoint(dirpath=...)` | [`bridge.training.config.CheckpointConfig.save`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.CheckpointConfig.save.html) | +| **Load directory** | `trainer.ckpt_path` | [`bridge.training.config.CheckpointConfig.load`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.CheckpointConfig.load.html) | +| **Pretrained checkpoint (for finetuning)** | `AutoResume.import_path` or manually load | [`bridge.training.config.CheckpointConfig.pretrained_checkpoint`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.CheckpointConfig.pretrained_checkpoint.html) | +| **Save frequency** | `ModelCheckpoint(every_n_train_steps=...)` | [`bridge.training.config.CheckpointConfig.save_interval`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.CheckpointConfig.save_interval.html) | +| **Save top-k** | `ModelCheckpoint(save_top_k=...)` | No direct equivalent - Megatron Bridge can keep the most recent checkpoints | +| **Most recent checkpoints** | No direct equivalent | [`bridge.training.config.CheckpointConfig.most_recent_k`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.CheckpointConfig.most_recent_k.html) | +| **Save last** | `ModelCheckpoint(save_last=...)` | Always enabled in Megatron Bridge | +| **Checkpoint format** | `strategy.save_ckpt_format` | [`bridge.training.config.CheckpointConfig.ckpt_format`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.CheckpointConfig.ckpt_format.html) | +| **Async saving** | `strategy.ckpt_async_save` | [`bridge.training.config.CheckpointConfig.async_save`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.CheckpointConfig.async_save.html) | +| **Parallel save** | `strategy.ckpt_parallel_save` | [`bridge.training.config.CheckpointConfig.fully_parallel_save`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.CheckpointConfig.fully_parallel_save.html) | +| **Parallel load** | `strategy.ckpt_parallel_load` | [`bridge.training.config.CheckpointConfig.fully_parallel_load`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.CheckpointConfig.fully_parallel_load.html) | +| **Load optimizer** | `strategy.ckpt_load_optimizer` | [`bridge.training.config.CheckpointConfig.load_optim`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.CheckpointConfig.load_optim.html) | +| **Save optimizer** | `strategy.ckpt_save_optimizer` | [`bridge.training.config.CheckpointConfig.save_optim`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.CheckpointConfig.save_optim.html) | +| **Load main params** | `strategy.ckpt_load_main_params` | [`bridge.training.config.CheckpointConfig.load_main_params_from_ckpt`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.CheckpointConfig.load_main_params_from_ckpt.html) | +| **Save weights only** | `ModelCheckpoint(save_weights_only=...)` | Inverse of `save_optim` | +| **Load strictness** | `strategy.ckpt_load_strictness` | [`bridge.training.config.CheckpointConfig.dist_ckpt_strictness`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.CheckpointConfig.dist_ckpt_strictness.html) | +| **Assume constant structure** | `strategy.ckpt_assume_constant_structure` | [`bridge.training.config.CheckpointConfig.ckpt_assume_constant_structure`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.CheckpointConfig.ckpt_assume_constant_structure.html) | +| **Save optim on train end** | `ModelCheckpoint(save_optim_on_train_end=...)` | Controlled by `save_optim` | +| **Resume from directory** | `AutoResume(resume_from_directory=...)` | [`bridge.training.config.CheckpointConfig.load`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.CheckpointConfig.load.html) | +| **Resume if exists** | `AutoResume(resume_if_exists=...)` | Automatic if `load` is set | +| **Resume ignore no checkpoint** | `AutoResume(resume_ignore_no_checkpoint=...)` | [`bridge.training.config.CheckpointConfig.exit_on_missing_checkpoint`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.CheckpointConfig.exit_on_missing_checkpoint.html) (inverse) | + +#### Before: NeMo 2.0 +```python +from nemo.lightning.pytorch.callbacks import ModelCheckpoint +from nemo.lightning import AutoResume, NeMoLogger + +# ModelCheckpoint callback +checkpoint_callback = ModelCheckpoint( + dirpath="/path/to/checkpoints", + every_n_train_steps=1000, + save_top_k=3, # Saves best 3 checkpoints based on monitored metric + save_last=True, + save_weights_only=False, + monitor="val_loss", # Metric to monitor for top-k selection +) + +# AutoResume for checkpoint resumption +resume = AutoResume( + resume_if_exists=True, + resume_ignore_no_checkpoint=True, + resume_from_directory="/path/to/checkpoints" +) + +# NeMoLogger ties everything together +logger = NeMoLogger( + log_dir="/path/to/logs", + name="my_experiment", + ckpt=checkpoint_callback, +) + +# MegatronStrategy parameters +strategy = run.Config( + MegatronStrategy, + save_ckpt_format="torch_dist", + ckpt_async_save=True, + ckpt_parallel_save=True, + ckpt_load_optimizer=True, + ckpt_save_optimizer=True, + ckpt_load_strictness=None, +) + +trainer = nl.Trainer(strategy=strategy) +logger.setup(trainer, resume.resume_if_exists) +resume.setup(trainer) +``` + +#### Now: Megatron Bridge +```python +checkpoint_config = CheckpointConfig( + # Saving configuration + save="/path/to/checkpoints", + save_interval=1000, + most_recent_k=3, # Keeps 3 most recent checkpoints (not metric-based) + save_optim=True, + save_rng=True, + + # Loading/resumption configuration + load="/path/to/checkpoints", # Resume from this directory (if exists) + load_optim=True, # Load optimizer state + exit_on_missing_checkpoint=False, # Don't exit if no checkpoint found (was resume_ignore_no_checkpoint) + + # Format and performance options + ckpt_format="torch_dist", + async_save=True, + fully_parallel_save=True, + fully_parallel_load=True, + dist_ckpt_strictness="assume_ok_unexpected", +) +``` + +**Key differences:** +- **Resume behavior**: Setting `load` enables automatic resume if checkpoint exists (no separate `AutoResume` needed) +- **Pretrained checkpoint**: Use `pretrained_checkpoint` to specify base model weights for fine-tuning (loaded before training starts) +- **Top-k**: NeMo 2.0's `save_top_k` monitors metrics; Megatron Bridge's `most_recent_k` keeps recent checkpoints +- **Configuration location**: All checkpoint settings unified in one config (not split across callback, logger, and strategy) + + +All checkpoint paths (`save`, `load`, `pretrained_checkpoint`) must point to **Megatron-format checkpoints**. Hugging Face checkpoints cannot be used directly—convert them first using [`bridge.models.conversion.auto_bridge.AutoBridge.import_ckpt`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.models.conversion.auto_bridge.AutoBridge.import_ckpt.html). See [Bridge Guide](/bridge-guide) for conversion details. + + +For comprehensive documentation on checkpoint formats, local checkpointing, fault tolerance, and advanced features, see [Checkpointing](/training/checkpointing). + +--- + +### Optimizer and LR Scheduler Migration + +Optimization configuration moves from NeMo 2.0's `MegatronOptimizerModule` approach to Megatron Bridge's direct [`megatron.core.optimizer.OptimizerConfig`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/megatron.core.optimizer.OptimizerConfig.html) and [`bridge.training.config.SchedulerConfig`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.SchedulerConfig.html). + +#### Before: NeMo 2.0 +```python +# NeMo 2.0 optimizer configuration with MegatronOptimizerModule +from nemo.lightning.pytorch.optim import MegatronOptimizerModule +from nemo.collections.llm.recipes import distributed_fused_adam_with_cosine_annealing + +# Option 1: Using recipe helper functions +optim_config = distributed_fused_adam_with_cosine_annealing( + max_lr=3e-4, + min_lr=3e-5, + warmup_steps=2000, +) + +# Option 2: Direct MegatronOptimizerModule +optim = MegatronOptimizerModule( + config=OptimizerConfig( + optimizer="adam", + lr=3e-4, + use_distributed_optimizer=True, + ), + lr_scheduler=CosineAnnealingScheduler( + warmup_steps=2000, + constant_steps=0, + decay_steps=100000, + ) +) +``` + +#### Now: Megatron Bridge +```python +# Megatron Bridge direct configuration +from megatron.bridge.recipes.utils.optimizer_utils import distributed_fused_adam_with_cosine_annealing + +# Option 1: Using utility functions +optimizer_config, scheduler_config = distributed_fused_adam_with_cosine_annealing( + max_lr=3e-4, + min_lr=3e-5, + lr_warmup_iters=2000, + lr_decay_iters=100000, +) + +# Option 2: Direct configuration +optimizer_config = OptimizerConfig( + optimizer="adam", + lr=3e-4, + min_lr=3e-5, + weight_decay=0.1, + use_distributed_optimizer=True, +) + +scheduler_config = SchedulerConfig( + lr_decay_style="cosine", + lr_warmup_iters=2000, + lr_decay_iters=100000, +) +``` + +### Logging Configuration Migration + +NeMo 2.0 uses `NeMoLogger` for TensorBoard and Weights & Biases (W&B) integration. Megatron Bridge consolidates logging configuration in [`bridge.training.config.LoggerConfig`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.LoggerConfig.html). + +#### Before: NeMo 2.0 + +```python +from nemo.lightning import NeMoLogger + +logger = NeMoLogger( + log_dir="/path/to/logs", + name="my_experiment", + use_datetime_version=True, + tensorboard=dict( + log_dir="/path/to/tensorboard", + ), + wandb=dict( + project="my_project", + name="my_run", + entity="my_team", + ), +) +``` + +#### Now: Megatron Bridge + +```python +from megatron.bridge.training.config import LoggerConfig + +logger_config = LoggerConfig( + # General logging + log_interval=10, # Log metrics every N iterations + log_throughput=True, # Log throughput per GPU + + # TensorBoard configuration + tensorboard_dir="/path/to/tensorboard", + tensorboard_log_interval=1, # Write to TensorBoard every N iterations + log_timers_to_tensorboard=False, + log_validation_ppl_to_tensorboard=False, + + # Weights & Biases configuration + wandb_project="my_project", + wandb_exp_name="my_run", + wandb_entity="my_team", + wandb_save_dir="/path/to/wandb", +) +``` + +**Key differences:** +- TensorBoard and W&B configuration unified in single `LoggerConfig` +- Fine-grained control over what gets logged (timers, memory, validation perplexity, etc.) +- No separate `NeMoLogger.setup()` call needed + +For more details on logging configuration and available options, see [Logging](/training/logging). + +### Profiling Configuration Migration + +Megatron Bridge centralizes all profiling functionality in [`bridge.training.config.ProfilingConfig`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.ProfilingConfig.html), replacing multiple NeMo callbacks. + +#### Nsys Profiling Migration + +##### Before: NeMo 2.0 +```python +# NeMo 2.0 used NsysCallback +from nemo.lightning.pytorch.callbacks import NsysCallback + +trainer = run.Config( + nl.Trainer, + callbacks=[ + NsysCallback( + start_step=100, + end_step=110, + ranks=[0], + gen_shape=True + ) + ] +) +``` + +##### Now: Megatron Bridge +```python +# Megatron Bridge uses ProfilingConfig +profiling_config = ProfilingConfig( + use_nsys_profiler=True, + profile_step_start=100, + profile_step_end=110, + profile_ranks=[0], + record_shapes=True, +) +``` + +#### PyTorch Profiler Migration + +##### Before: NeMo 2.0 +```python +# NeMo 2.0 used PytorchProfilerCallback +from nemo.lightning.pytorch.callbacks import PytorchProfilerCallback + +trainer = run.Config( + nl.Trainer, + callbacks=[ + PytorchProfilerCallback( + start_step=100, + end_step=110, + warmup_steps=1, + active_steps=5, + trace_dir="/path/to/traces", + ) + ] +) +``` + +##### Now: Megatron Bridge +```python +# Megatron Bridge uses ProfilingConfig +profiling_config = ProfilingConfig( + use_pytorch_profiler=True, + profile_step_start=100, + profile_step_end=110, + profile_ranks=[0], + record_memory_history=True, + memory_snapshot_path="memory_profile.pickle", +) +``` + +### PEFT Configuration Migration + +PEFT (Parameter-Efficient Fine-Tuning) enables fine-tuning with a small fraction of trainable parameters by freezing the base model and training only adapter modules. + +#### Before: NeMo 2.0 + +```python +from nemo.collections import llm +import nemo_run as run + +# Create PEFT configuration +lora = llm.peft.LoRA( + target_modules=['linear_qkv', 'linear_proj'], + dim=32, + alpha=16, + dropout=0.0, +) + +# Pass to finetune() +llm.finetune( + model=model, + data=data, + trainer=trainer, + peft=lora, # PEFT as argument +) +``` + +#### Now: Megatron Bridge + +```python +from megatron.bridge.peft import LoRA +from megatron.bridge.training.config import ConfigContainer, CheckpointConfig + +# Include PEFT in ConfigContainer +config = ConfigContainer( + model=Llama3ModelProvider8B(), + # ... other configs + checkpoint=CheckpointConfig( + pretrained_checkpoint="/path/to/megatron/checkpoint", # Required for PEFT + save="/path/to/peft/checkpoints", + ), + peft=LoRA( + target_modules=["linear_qkv", "linear_proj", "linear_fc1", "linear_fc2"], + dim=32, + alpha=16, + dropout=0.0, + ), +) +``` + +**Key differences:** +- PEFT config is part of `ConfigContainer`, not a separate argument to `finetune()` +- Must set `checkpoint.pretrained_checkpoint` when using PEFT (enforced at validation) +- Target module names are the same between NeMo 2.0 and Megatron Bridge + +**Supported PEFT methods:** +- **LoRA**: Low-Rank Adaptation via [`bridge.peft.lora.LoRA`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.peft.lora.LoRA.html) +- **DoRA**: Weight-Decomposed Low-Rank Adaptation via [`bridge.peft.dora.DoRA`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.peft.dora.DoRA.html) + +For comprehensive PEFT documentation including adapter design, checkpoint handling, wildcard targeting, and best practices, see [Peft](/training/peft). + +## Entry Points: `pretrain` and `finetune` + +NeMo 2.0's `llm.pretrain()` and `llm.finetune()` API functions map directly to Megatron Bridge's entry point functions with unified configuration. + +### NeMo 2.0 Entry Points + +In NeMo 2.0, you call `llm.pretrain()` or `llm.finetune()` from `nemo.collections.llm.api`: + +```python +from nemo.collections import llm +import nemo_run as run + +# Pretraining +result = llm.pretrain( + model=model_config, + data=data_config, + trainer=trainer_config, + log=logger_config, + resume=resume_config, + optim=optimizer_config, +) + +# Fine-tuning +result = llm.finetune( + model=model_config, + data=data_config, + trainer=trainer_config, + log=logger_config, + resume=resume_config, + optim=optimizer_config, + peft=peft_config, # Optional PEFT + tokenizer="model", # or "data" +) +``` + +### Megatron Bridge Entry Points + +In Megatron Bridge, training entry points take a single `ConfigContainer` and a `forward_step_func`: + +```python +from megatron.bridge.training import pretrain, finetune +from megatron.bridge.training.config import ConfigContainer + +# Create unified configuration +cfg = ConfigContainer( + model=model_provider, + train=train_config, + dataset=dataset_config, + optimizer=optimizer_config, + scheduler=scheduler_config, + checkpoint=checkpoint_config, + logger=logger_config, + mixed_precision="bf16_mixed", + # peft=peft_config, # Optional for fine-tuning +) + +# Pretraining +from megatron.bridge.training.gpt_step import forward_step +pretrain(cfg, forward_step_func=forward_step) + +# Fine-tuning (same function signature) +finetune(cfg, forward_step_func=forward_step) +``` + +#### Understanding `forward_step_func` + +The `forward_step_func` combines three responsibilities into a single function: + +1. **Fetch a batch** from the data iterator +2. **Run the forward pass** through the model +3. **Define the loss function** to compute loss from the model output + +**Signature:** +```python +def forward_step( + state: GlobalState, + data_iterator: Iterable, + model: nn.Module, +) -> tuple[torch.Tensor, Callable]: + """ + Args: + state: Global training state (contains config, timers, etc.) + data_iterator: Iterator over training/validation data + model: The model to run forward pass on + + Returns: + output_tensor: Model output (logits) + loss_func: Callable that computes loss from output_tensor + """ +``` + +For GPT models, use the provided [`bridge.training.gpt_step.forward_step`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.gpt_step.forward_step.html). For custom models or specialized training logic, implement your own following this pattern. + +**Key differences:** +- All configuration consolidated into single `ConfigContainer` object +- Training mode determined by dataset type and checkpoint configuration, not separate function calls +- Must provide `forward_step_func` that handles batch fetching, forward pass, and loss computation +- No separate `resume`, `log`, `optim` arguments - all configurations are part of the `ConfigContainer` + +### `pretrain` + +Use `pretrain()` with `GPTDatasetConfig` for training models from scratch: + +```python +from megatron.bridge.training import pretrain +from megatron.bridge.training.gpt_step import forward_step + +config = ConfigContainer( + model=Llama3ModelProvider8B( + tensor_model_parallel_size=2, + pipeline_model_parallel_size=2, + ), + train=TrainingConfig( + train_iters=100000, + eval_interval=1000, + micro_batch_size=1, + global_batch_size=512, + ), + dataset=GPTDatasetConfig( + blend=["/path/to/train_data_text_document"], + seq_length=4096, + split="949,50,1", + ), + optimizer=OptimizerConfig(optimizer="adam", lr=3e-4), + checkpoint=CheckpointConfig(save="/path/to/checkpoints", save_interval=1000), + mixed_precision="bf16_mixed", +) + +pretrain(config, forward_step_func=forward_step) +``` + +### `finetune` + +Use `finetune()` with `FinetuningDatasetConfig` for both full fine-tuning (SFT) and parameter-efficient fine-tuning (PEFT): + +#### Supervised Fine-Tuning (SFT) + +Full fine-tuning without PEFT - all model parameters are updated: + +```python +from megatron.bridge.training import finetune +from megatron.bridge.training.gpt_step import forward_step + +config = ConfigContainer( + model=Llama3ModelProvider8B(), + train=TrainingConfig( + train_iters=1000, + eval_interval=100, + micro_batch_size=1, + global_batch_size=128, + ), + dataset=FinetuningDatasetConfig( + dataset_root="/path/to/instruction_data", + seq_length=4096, + do_validation=True, + ), + checkpoint=CheckpointConfig( + pretrained_checkpoint="/path/to/megatron/checkpoint", # Must be Megatron format + save="/path/to/sft_checkpoints", + ), + optimizer=OptimizerConfig(optimizer="adam", lr=1e-5), + mixed_precision="bf16_mixed", +) + +finetune(config, forward_step_func=forward_step) +``` + +#### Parameter-Efficient Fine-Tuning (PEFT) + +Add a `peft` configuration to enable parameter-efficient training: + +```python +from megatron.bridge.peft import LoRA + +config = ConfigContainer( + model=Llama3ModelProvider8B(), + train=TrainingConfig( + train_iters=1000, + eval_interval=100, + micro_batch_size=1, + global_batch_size=128, + ), + dataset=FinetuningDatasetConfig( + dataset_root="/path/to/instruction_data", + seq_length=4096, + do_validation=True, + ), + checkpoint=CheckpointConfig( + pretrained_checkpoint="/path/to/megatron/checkpoint", + save="/path/to/peft_checkpoints", + ), + peft=LoRA( + target_modules=["linear_qkv", "linear_proj", "linear_fc1", "linear_fc2"], + dim=32, + alpha=16, + ), + optimizer=OptimizerConfig(optimizer="adam", lr=1e-4), + mixed_precision="bf16_mixed", +) + +finetune(config, forward_step_func=forward_step) +``` + +**Converting Hugging Face checkpoints**: If you have a Hugging Face model, convert it to Megatron checkpoint format first: + +```python +from megatron.bridge import AutoBridge + +# Convert HuggingFace to Megatron format +AutoBridge.import_ckpt( + "meta-llama/Meta-Llama-3-8B", + "/path/to/megatron/checkpoint" +) +``` + +See [Bridge Guide](/bridge-guide) for more details on model conversion. + +### Advanced: Custom Forward Step and Loss Reduction + +For comprehensive documentation on entry points, forward step functions, and customization patterns, see [Entry Points](/training/entry-points). + +#### Forward Step Customization + +In NeMo 2.0, custom `forward_step` and `data_step` functions can be attached to the model configuration. In Megatron Bridge, the forward step function is passed directly as an argument to `pretrain()` or `finetune()`. + +##### NeMo 2.0: Custom Steps Attached to Config + +```python +# NeMo 2.0: Define custom functions and attach to model config +import torch + +def custom_forward_step(model, batch) -> torch.Tensor: + """Custom forward step for specialized loss computation.""" + output = model(batch['tokens'], batch['attention_mask']) + loss = compute_custom_loss(output, batch['labels']) + return loss + +# Attach to config in NeMo 2.0 +model_config = llm.Llama3Config8B() +model_config.forward_step_fn = custom_forward_step # Override default forward step + +model = run.Config(llm.LlamaModel, config=model_config) +``` + +##### Megatron Bridge: Pass Custom Forward Step + +```python +# Megatron Bridge: Define and pass forward step function +import torch +from typing import Iterable +from functools import partial +from megatron.bridge.training.state import GlobalState + +def custom_forward_step( + state: GlobalState, + data_iterator: Iterable, + model: torch.nn.Module, +) -> tuple[torch.Tensor, partial]: + """Custom forward step for specialized loss computation.""" + # Get batch from iterator + batch = next(data_iterator) + tokens = batch['tokens'].cuda() + labels = batch['labels'].cuda() + loss_mask = batch['loss_mask'].cuda() + + # Custom forward logic + output = model(tokens, attention_mask=batch.get('attention_mask')) + + # Define custom loss function + def loss_func(output_tensor): + return compute_custom_loss(output_tensor, labels, loss_mask) + + return output, loss_func + +# Pass to training function +pretrain(cfg, forward_step_func=custom_forward_step) +``` + +#### Loss Reduction Pattern + +NeMo 2.0 uses `MegatronLossReduction` for custom loss computation and reduction across microbatches. Megatron Bridge achieves the same through the loss function returned by `forward_step`. + +##### NeMo 2.0: MegatronLossReduction + +```python +from nemo.lightning.megatron_parallel import MegatronLossReduction + +class CustomLossReduction(MegatronLossReduction): + def forward(self, batch, forward_out): + """Compute loss from forward output.""" + loss = compute_loss(forward_out, batch['labels']) + return loss, {"custom_metric": some_metric} + + def reduce(self, losses_reduced_per_micro_batch): + """Reduce losses across microbatches.""" + losses = [x["custom_metric"] for x in losses_reduced_per_micro_batch] + return torch.stack(losses).mean() + +# Attach to model +model._training_loss_reduction = CustomLossReduction() +``` + +##### Megatron Bridge: Loss Function Pattern + +```python +def custom_forward_step(state, data_iterator, model): + """Forward step with custom loss reduction.""" + batch = next(data_iterator) + tokens = batch['tokens'].cuda() + labels = batch['labels'].cuda() + loss_mask = batch['loss_mask'].cuda() + + output = model(tokens) + + def loss_func(output_tensor): + """Compute and return loss in reduction-friendly format. + + Return formats: + - Single value: loss (averaged over microbatches only) + - Tuple: (loss, num_tokens) - averaged over microbatches and tokens + - Dict: {"loss": loss, "custom_metric": value, ...} - for logging + """ + loss = compute_loss(output_tensor, labels, loss_mask) + num_tokens = loss_mask.sum() + + # Return (loss, num_tokens) for proper averaging + # Training loop automatically reduces across microbatches and data parallel ranks + return { + "loss": torch.cat([loss.view(1), num_tokens.view(1)]), + "custom_metric": torch.cat([some_metric.view(1), num_tokens.view(1)]), + } + + return output, loss_func + +# Pass to training - reduction handled automatically +pretrain(cfg, forward_step_func=custom_forward_step) +``` + +**Key differences:** +- **NeMo 2.0**: Separate `MegatronLossReduction` class with `forward()` and `reduce()` methods +- **Megatron Bridge**: Loss function returns dict with format `{key: [value, count]}` for automatic reduction +- **Reduction logic**: Megatron Bridge automatically averages `value/count` across microbatches and data parallel ranks + +The training loop in Megatron Bridge (see [`bridge.training.train.train_step`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.train.train_step.html)) automatically: +1. Calls the loss function for each microbatch +2. Aggregates results across microbatches +3. Performs data-parallel all-reduce +4. Computes final averaged values + +#### When to Customize + +Use custom forward steps when you need: +- Custom loss functions beyond standard language modeling +- Multi-task learning with multiple loss components +- Additional metrics computed during training +- Specialized batch preprocessing + +For complete documentation on entry point signatures, loss calculation patterns, state access, and more advanced customization options, see [Entry Points](/training/entry-points). + +--- + +## Callback Migration + +Megatron Bridge converts most NeMo 2.0 callbacks into explicit configuration fields or utility functions. + +### DDP Parity Checker + +Validates that model weights are synchronized across data-parallel replicas. + +#### Before: NeMo 2.0 +```python +from nemo.lightning.pytorch.callbacks import DDPParityChecker + +trainer = run.Config( + nl.Trainer, + callbacks=[DDPParityChecker(check_interval=100)] +) +``` + +#### Now: Megatron Bridge +```python +# Built into TrainingConfig +train_config = TrainingConfig( + check_weight_hash_across_dp_replicas_interval=100, +) +``` + +### Garbage Collection + +Manual garbage collection to free memory during training. + +#### Before: NeMo 2.0 +```python +from nemo.lightning.pytorch.callbacks import GarbageCollectionCallback + +trainer = run.Config( + nl.Trainer, + callbacks=[ + GarbageCollectionCallback( + gc_interval_train=100, + gc_interval_val=100, + ) + ] +) +``` + +#### Now: Megatron Bridge +```python +# Built into TrainingConfig +train_config = TrainingConfig( + manual_gc=True, # Enable manual garbage collection + manual_gc_interval=100, # GC interval during training (was gc_interval_train) + manual_gc_eval=True, # Enable GC at start/end of evaluation (was gc_interval_val) +) +``` + +### Communication Overlap + +Enables overlapping of tensor/pipeline parallel communication with computation. + +#### Before: NeMo 2.0 +```python +from nemo.lightning.pytorch.callbacks import MegatronCommOverlapCallback + +trainer = run.Config( + nl.Trainer, + callbacks=[ + MegatronCommOverlapCallback( + tp_comm_overlap=True, + ... + ) + ] +) +``` + +#### Now: Megatron Bridge +```python +from megatron.bridge.training.comm_overlap import CommOverlapConfig + +config = ConfigContainer( + comm_overlap=CommOverlapConfig( + tp_comm_overlap=True, + tp_comm_overlap_cfg=..., # Detailed TP overlap settings + ), +) +``` + +For comprehensive documentation on communication overlap strategies (TP, PP, DP, CP, MoE), hardware requirements, and performance tuning, see [Communication Overlap](/training/communication-overlap). + +### Preemption Handling + +Gracefully handle SLURM/cluster preemption signals. + +#### Before: NeMo 2.0 +```python +from nemo.lightning.pytorch.callbacks import PreemptionCallback + +trainer = run.Config( + nl.Trainer, + callbacks=[PreemptionCallback()] +) +``` + +#### Now: Megatron Bridge +```python +# Built into TrainingConfig +train_config = TrainingConfig( + exit_signal_handler=True, # Enable preemption signal handling +) +``` + +For more details on preemption handling and fault tolerance, see [Resiliency](/training/resiliency). + +### Experimental Features + +Enable Megatron Core experimental features. + +#### Before: NeMo 2.0 +```python +from nemo.lightning.pytorch.callbacks import MegatronEnableExperimentalCallback + +trainer = run.Config( + nl.Trainer, + callbacks=[MegatronEnableExperimentalCallback()] +) +``` + +#### Now: Megatron Bridge +```python +from megatron.bridge.training.config import DistributedInitConfig + +dist_config = DistributedInitConfig( + enable_megatron_core_experimental=True, +) +``` + +### MoE Token Drop + +Configures expert capacity and token padding for MoE models. + +#### Before: NeMo 2.0 +```python +from nemo.lightning.pytorch.callbacks import MegatronTokenDropCallback + +callbacks = [ + MegatronTokenDropCallback( + moe_expert_capacity_factor=1.0, + moe_pad_expert_input_to_capacity=True + ) +] +``` + +#### Now: Megatron Bridge +```python +from megatron.bridge.training.utils.moe_token_drop import apply_moe_token_drop + +model = GPTModelProvider( + # MoE architecture + num_moe_experts=8, + moe_router_topk=2, + moe_token_dispatcher_type="alltoall", +) + +# Apply token drop optimization +apply_moe_token_drop( + model, + moe_expert_capacity_factor=1.0, + moe_pad_expert_input_to_capacity=True +) +``` + +### DeepEP for MoE + +Enables DeepEP optimizations for MoE models on supported hardware (Ampere/Hopper GPUs). + +#### Before: NeMo 2.0 +```python +from nemo.lightning.pytorch.callbacks import DeepEPCallback + +callbacks = [DeepEPCallback()] # Automatically applies if hardware supports it +``` + +#### Now: Megatron Bridge +```python +from megatron.bridge.training.deepep import apply_deepep + +model = GPTModelProvider( + num_moe_experts=8, + # ... other MoE settings +) + +# Apply DeepEP optimizations (only on Ampere/Hopper GPUs) +# Hardware validation is performed automatically during training +apply_deepep(model) +``` + +--- + +## NeMo-Run, Plugins, and Launching + +Megatron Bridge supports both direct Python execution and NeMo-Run orchestration. While NeMo 2.0 relied heavily on NeMo-Run's recipe system, Megatron Bridge provides more flexibility. + +For complete details on launching training jobs, configuration overrides, and NeMo-Run integration, see [Recipe Usage](/recipe-usage). + +### NeMo-Run Integration + +#### Direct Python Execution +Megatron Bridge supports standard PyTorch distributed execution patterns: + +```bash +# Direct script execution with torchrun +python -m torch.distributed.run --nproc_per_node=8 my_training_script.py + +# Multi-node execution +torchrun --nnodes=4 --nproc_per_node=8 \ + --master_addr="node0" --master_port=12345 \ + my_training_script.py +``` + +#### NeMo-Run with Plugins (Script Mode Recommended) +If using NeMo-Run, **strongly recommend using `run.Script` mode** for better dependency management. Megatron Bridge plugins are designed to work well with this approach: + +```python +# Megatron Bridge NeMo-Run integration +import nemo_run as run +from megatron.bridge.recipes.run_plugins import ( + NsysPlugin, WandbPlugin, PreemptionPlugin, FaultTolerancePlugin +) + +# Create task +task = run.Script("my_training_script.py", args=[]) + +# Configure executor with plugins +executor = run.SlurmExecutor(nodes=2, nproc_per_node=8) +executor.plugins = [ + NsysPlugin(profile_step_start=100, profile_step_end=110), + WandbPlugin(project="my_project", entity="my_team"), + PreemptionPlugin(preempt_time=120), +] + +# Submit job +run.run(task, executor=executor) +``` + +### Plugin Migration Comparison + +| **Plugin** | **NeMo 2.0** | **Megatron Bridge** | +|------------|-------------|-----------| +| **Nsys** | `nemo.lightning.run.plugins.NsysPlugin` | [`bridge.recipes.run_plugins.NsysPlugin`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.recipes.run_plugins.NsysPlugin.html) | +| **Wandb** | `nemo.lightning.run.plugins.WandbPlugin` | [`bridge.recipes.run_plugins.WandbPlugin`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.recipes.run_plugins.WandbPlugin.html) | +| **Preemption** | `nemo.lightning.run.plugins.PreemptionPlugin` | [`bridge.recipes.run_plugins.PreemptionPlugin`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.recipes.run_plugins.PreemptionPlugin.html) | +| **Fault Tolerance** | `nemo.lightning.run.plugins.FaultTolerancePlugin` | [`bridge.recipes.run_plugins.FaultTolerancePlugin`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.recipes.run_plugins.FaultTolerancePlugin.html) | diff --git a/fern/v0.2.0/pages/parallelisms.mdx b/fern/v0.2.0/pages/parallelisms.mdx new file mode 100644 index 0000000000..682c173faf --- /dev/null +++ b/fern/v0.2.0/pages/parallelisms.mdx @@ -0,0 +1,384 @@ +--- +title: "Parallelisms Guide" +description: "" +--- + +Megatron Bridge supports various data-parallel and model-parallel deep learning workload deployment methods, which can be mixed together arbitrarily. These parallelism strategies are configured through model provider classes and leverage Megatron Core's implementation for performance and memory efficiency. + +## Data Parallelism + +Data Parallelism (DP) replicates the model across multiple GPUs. Data batches are evenly distributed between GPUs and the data-parallel GPUs process them independently. While the computation workload is efficiently distributed across GPUs, inter-GPU communication is required to keep the model replicas consistent between training steps. + +### Distributed Data Parallelism + +Distributed Data Parallelism (DDP) keeps the model copies consistent by synchronizing parameter gradients across data-parallel GPUs before each parameter update. More specifically, it sums the gradients of all model copies using all-reduce communication collectives. + +![Distributed Data Parallelism](/assets/training/images/ddp.gif) +*Figure: Distributed Data Parallelism synchronizes gradients across multiple GPUs using all-reduce operations.* + +### Distributed Optimizer + +[Distributed optimizer](https://docs.nvidia.com/megatron-core/developer-guide/latest/user-guide/features/dist_optimizer.html) is a memory-optimized data-parallel deployment method. It shards the optimizer states and the high-precision master parameters across data-parallel GPUs instead of replicating them. At the parameter optimizer step, each data-parallel GPU updates its shard of parameters. Since each GPU needs its own gradient shard, the distributed optimizer conducts reduce-scatter of the parameter gradients instead of all-reduce of them. Then, the updated parameter shards are all-gathered across data-parallel GPUs. This approach significantly reduces the memory need of large-scale LLM training. + +### Enable Data Parallelism + +In Megatron Bridge, DDP is the default parallel deployment method. The total number of GPUs corresponds to the size of the DP group, and training an LLM with model parallelism decreases the size of the DP group. + +To enable the distributed optimizer, configure the [`bridge.training.config.OptimizerConfig`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.OptimizerConfig.html) and [`bridge.training.config.DistributedDataParallelConfig`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.DistributedDataParallelConfig.html) + +```python +from megatron.bridge.training.config import ConfigContainer, DistributedDataParallelConfig, OptimizerConfig + +optimizer_config = OptimizerConfig( + optimizer="adam", + lr=3e-4, + weight_decay=0.1, + adam_beta1=0.9, + adam_beta2=0.95, + use_distributed_optimizer=True, + clip_grad=1.0, +) +ddp_config = DistributedDataParallelConfig(use_distributed_optimizer=True) + +config = ConfigContainer( + ddp=ddp_config, + optimizer=optimizer_config, + # ... other config parameters +) +``` + +For more optimizer options, refer to the [`bridge.training.config.OptimizerConfig`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.OptimizerConfig.html) API documentation. + +## Model Parallelism + +Model Parallelism (MP) is a distributed model deployment method that partitions the model parameters across GPUs to reduce the need for per-GPU memory. Megatron Bridge supports various model-parallel methods through Megatron Core, which can be mixed to maximize LLM training performance. + +### Tensor Parallelism + +Tensor Parallelism (TP) is a model-parallel partitioning method that distributes the parameter tensor of an individual layer across GPUs. In addition to reducing model state memory usage, it also saves activation memory as the per-GPU tensor sizes shrink. However, the reduced per-GPU tensor size increases CPU overhead due to smaller per-GPU kernel workloads. + +![Tensor Parallelism Overview](/assets/training/images/tp1.png) +*Figure 1: Tensor Parallelism distributes individual layer parameters across multiple GPUs.* + +![Tensor Parallelism Implementation](/assets/training/images/tp2.png) +*Figure 2: Detailed view of how tensor parallelism splits weight matrices and synchronizes computations.* + +#### Enable Tensor Parallelism + +To enable TP in Megatron Bridge, configure the `tensor_model_parallel_size` parameter in your model provider. This parameter determines the number of GPUs among which the model's tensors are partitioned. + +```python +from megatron.bridge.models import GPTModelProvider +from megatron.bridge.training.config import ConfigContainer + +# Configure model with tensor parallelism +model_config = GPTModelProvider( + tensor_model_parallel_size=2, # Enable TP across 2 GPUs + # ... other model parameters +) + +config = ConfigContainer( + model=model_config, + # ... other config parameters +) +``` + +#### Implement Tensor Parallelism + +Megatron Bridge integrates TP through the implementation from Megatron Core. For detailed API usage and additional configurations, consult the [Megatron Core Developer Guide](https://docs.nvidia.com/megatron-core/developer-guide/latest/apidocs/core/core.tensor_parallel.html). + +### Pipeline Parallelism + +Pipeline Parallelism (PP) is a technique that assigns consecutive layers or segments of a neural network to different GPUs. This division allows each GPU to process different stages of the network sequentially. + +![Pipeline Parallelism](/assets/training/images/pp.gif) +*Figure: Pipeline Parallelism distributes consecutive layers across multiple GPUs, processing batches in a pipeline fashion.* + +#### Enable Pipeline Parallelism + +To utilize Pipeline Parallelism in Megatron Bridge, set the `pipeline_model_parallel_size` parameter in your model configuration. This parameter specifies the number of GPUs among which the model's layers are distributed. + +```python +from megatron.bridge.models import GPTModelProvider +from megatron.bridge.training.config import ConfigContainer + +# Configure model with pipeline parallelism +model_config = GPTModelProvider( + pipeline_model_parallel_size=4, # Distribute layers across 4 GPUs + # ... other model parameters +) + +config = ConfigContainer( + model=model_config, + # ... other config parameters +) +``` + +#### Interleaved Pipeline Parallel Schedule + +To minimize the pipeline bubble, the computation on each GPU can be divided into multiple subsets of layers (referred to as model chunks), rather than a single contiguous block. Enable this by setting `virtual_pipeline_model_parallel_size`: + +```python +model_config = GPTModelProvider( + pipeline_model_parallel_size=4, + virtual_pipeline_model_parallel_size=2, # 2 model chunks per pipeline stage + # ... other model parameters +) +``` + +For more insights into this approach, see the detailed blog: [Scaling Language Model Training](https://developer.nvidia.com/blog/scaling-language-model-training-to-a-trillion-parameters-using-megatron/#pipeline_parallelism). + +#### Implement Pipeline Parallelism + +The Megatron Bridge implementation of PP leverages functionalities from Megatron Core. For more detailed API usage and configurations related to PP, visit the [Megatron Core Developer Guide](https://docs.nvidia.com/megatron-core/developer-guide/latest/apidocs/core/core.pipeline_parallel.html). + +### Expert Parallelism and Mixture of Experts (MoE) + +Expert Parallelism (EP) is a type of model parallelism that distributes experts of a Mixture of Experts (MoE) model across GPUs. Unlike other model-parallel techniques, EP is applied to only the expert layers and does not impact the parallel mapping of the rest of the layers. + +MoE is a machine learning technique where multiple specialized models (experts, usually multi-layer perceptrons) are combined to solve a complex task. Each expert focuses on a specific subtask or domain, while a gating network dynamically activates the most appropriate expert based on the current input. + +![Expert Parallelism](/assets/training/images/ep.png) +*Figure: Expert Parallelism distributes MoE experts across multiple GPUs while keeping other layers replicated.* + +#### Basic MoE Configuration + +To enable MoE in Megatron Bridge, configure the basic MoE parameters in your model provider: + +```python +from megatron.bridge.models import GPTModelProvider + +# Configure basic MoE model +model_config = GPTModelProvider( + num_moe_experts=8, # Number of experts in the MoE module + moe_router_topk=2, # Number of experts activated per token + moe_ffn_hidden_size=8192, # Hidden size for expert FFN layers + # ... other model parameters +) +``` + +#### Enable Expert Parallelism + +To enable EP, set `expert_model_parallel_size` in your model configuration. For example, if the model has eight experts (`num_moe_experts=8`), then setting `expert_model_parallel_size=4` results in each GPU processing two experts. The number of experts should be divisible by the expert parallel size. + +```python +# Configure MoE model with expert parallelism +model_config = GPTModelProvider( + num_moe_experts=8, + expert_model_parallel_size=4, # Distribute 8 experts across 4 GPUs (2 experts per GPU) + # ... other model parameters +) +``` + +#### Enable Expert Tensor Parallelism + +To enable Expert Tensor Parallelism (ETP), set `expert_tensor_parallel_size` in your model configuration: + +```python +model_config = GPTModelProvider( + num_moe_experts=8, + expert_model_parallel_size=4, + expert_tensor_parallel_size=2, # Apply tensor parallelism within each expert + # ... other model parameters +) +``` + +#### Advanced MoE Features + +##### DeepEP Optimization + +Megatron Bridge includes DeepEP optimization for improved MoE performance on Ampere and Hopper GPUs: + +```python +from megatron.bridge.training.deepep import apply_deepep + +# Apply DeepEP optimization +apply_deepep(model_config) # Sets moe_token_dispatcher_type="flex" and other optimizations +``` + +##### Token Dropping for Load Balancing + +Token dropping improves performance by balancing work across experts: + +```python +from megatron.bridge.training.utils.moe_token_drop import apply_moe_token_drop + +# Apply token drop settings for load balancing +apply_moe_token_drop(model_config, moe_expert_capacity_factor=1.0) +``` + +#### Complete MoE Configuration Example + +```python +from megatron.bridge.models import GPTModelProvider +from megatron.bridge.training.config import ConfigContainer +from megatron.bridge.training.deepep import apply_deepep +from megatron.bridge.training.utils.moe_token_drop import apply_moe_token_drop + +# Configure MoE model with expert parallelism +model_config = GPTModelProvider( + num_layers=32, + hidden_size=4096, + num_attention_heads=32, + + # MoE configuration + num_moe_experts=8, # 8 experts total + moe_router_topk=2, # Activate 2 experts per token + + # Expert parallelism + expert_model_parallel_size=4, # Distribute experts across 4 GPUs + expert_tensor_parallel_size=2, # Apply TP within each expert + + # ... other model parameters +) + +config = ConfigContainer( + model=model_config, + # ... other config parameters +) +``` + +#### Expert Parallelism Implementation + +The Megatron Bridge implementation of EP uses functionality from Megatron Core. Please consult the [Megatron Core MoE layer](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/transformer/moe/moe_layer.py#L42) for more MoE implementation details. + +## Activation Partitioning + +In LLM training, a large memory space is needed to store the input activations of the network layers. Megatron Bridge provides effective activation distribution methods through Megatron Core, which is critical in training LLMs with large sequence lengths or large per-GPU micro-batch sizes. + +### Sequence Parallelism + +Sequence Parallelism (SP) extends tensor-level model parallelism by distributing computing load and activation memory across multiple GPUs along the sequence dimension of transformer layers. This method is particularly useful for portions of the layer that have previously not been parallelized, enhancing overall model performance and efficiency. + +![Sequence Parallelism](/assets/training/images/sp.png) +*Figure: Sequence Parallelism distributes the sequence dimension across multiple GPUs, reducing activation memory.* + +#### Enable Sequence Parallelism + +To utilize SP in Megatron Bridge, set the `sequence_parallel` parameter to `True` in your model configuration. Note that this feature is effective only when the tensor parallel size (`tensor_model_parallel_size`) is greater than `1`. + +```python +from megatron.bridge.models import GPTModelProvider + +# Configure model with sequence parallelism +model_config = GPTModelProvider( + tensor_model_parallel_size=2, # Required for sequence parallelism + sequence_parallel=True, # Enable sequence parallelism + # ... other model parameters +) +``` + +#### Implement Sequence Parallelism + +The Megatron Bridge implementation of SP utilizes functionality from Megatron Core. For an in-depth look at how Sequence Parallelism is integrated into the Megatron Core architecture, you can examine the source code: [Megatron-LM Sequence Parallel Source Code](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/tensor_parallel/layers.py). + +### Context Parallelism + +Context Parallelism (CP) is a method for parallelizing the processing of neural network activations across multiple GPUs by partitioning the input tensors along the sequence dimension. Unlike Sequence Parallelism (SP) that partitions the activations of specific layers, CP divides the activations of all layers. + +CP is critical for training long context models, as it allows the model to handle longer sequences by distributing the sequence activations across multiple GPUs. This method reduces the memory footprint and computational cost of processing long sequences. + +#### Enable Context Parallelism + +To activate CP in Megatron Bridge, set the `context_parallel_size` parameter in your model configuration. This parameter specifies the number of GPUs across which the model's sequence activations are distributed. + +```python +from megatron.bridge.models import GPTModelProvider + +# Configure model with context parallelism +model_config = GPTModelProvider( + context_parallel_size=2, # Distribute sequence across 2 GPUs + # ... other model parameters +) +``` + +For long context training scenarios, context parallelism is particularly effective and essential for handling sequences that exceed the memory capacity of individual GPUs. + +#### Implement Context Parallelism + +Megatron Bridge leverages functionalities from both Megatron Core and Transformer Engine to implement CP efficiently. During forward propagation, each GPU handles a segment of the sequence, storing only the necessary Key and Value (KV) pairs. In the backward pass, these KV pairs are reassembled across GPUs using advanced communication schemes like all-gather and reduce-scatter transformed into point-to-point communications in a ring topology. This method reduces the memory footprint significantly while maintaining computational efficiency. + +For more detailed technical information and implementation details, visit: +- [Megatron Core Context Parallelism Documentation](https://docs.nvidia.com/megatron-core/developer-guide/latest/user-guide/features/context_parallel.html) +- [Megatron Core wrappers for Transformer Engine](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/transformer/custom_layers/transformer_engine.py) +- [Transformer Engine attention modules](https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/pytorch/attention.py) + +## Combined Parallelism Example + +Megatron Bridge allows you to combine multiple parallelism strategies for optimal performance and memory efficiency: + +```python +from megatron.bridge.models import GPTModelProvider +from megatron.bridge.training.config import ConfigContainer, OptimizerConfig + +# Configure model with multiple parallelism strategies +model_config = GPTModelProvider( + # Model parallelism + tensor_model_parallel_size=2, # 2-way tensor parallelism + pipeline_model_parallel_size=4, # 4-way pipeline parallelism + virtual_pipeline_model_parallel_size=2, # Interleaved pipeline + + # Activation partitioning + sequence_parallel=True, # Enable sequence parallelism (requires TP > 1) + context_parallel_size=2, # 2-way context parallelism + + # Expert parallelism (for MoE models) + num_moe_experts=8, # 8 experts + expert_model_parallel_size=4, # Distribute experts across 4 GPUs + + # ... other model parameters +) + +# Configure distributed optimizer +optimizer_config = OptimizerConfig( + optimizer="adam", + use_distributed_optimizer=True, # Enable distributed optimizer + # ... other optimizer parameters +) + +config = ConfigContainer( + model=model_config, + optimizer=optimizer_config, + # ... other config parameters +) +``` + +## Data Parallel Size Calculation + +The data parallel size is automatically calculated based on the total world size and model parallelism settings: + +``` +data_parallel_size = world_size / (tensor_model_parallel_size × pipeline_model_parallel_size × context_parallel_size) +``` + +For example, with 32 GPUs total and the configuration above: +- `tensor_model_parallel_size = 2` +- `pipeline_model_parallel_size = 4` +- `context_parallel_size = 2` +- `data_parallel_size = 32 / (2 × 4 × 2) = 2` + +## Configuration Guidelines + +### Memory Optimization +- Use **distributed optimizer** to reduce optimizer state memory +- Enable **sequence parallelism** when using tensor parallelism to reduce activation memory +- Use **context parallelism** for long sequence training +- Consider **pipeline parallelism** for very large models that don't fit on a single GPU + +### Performance Optimization +- **Tensor parallelism** works best within a single node (high bandwidth) +- **Pipeline parallelism** can work across nodes but requires careful batch size tuning +- **Context parallelism** is essential for long context scenarios +- **Expert parallelism** is specific to MoE models and should match the number of experts + +### Compatibility +- **Sequence parallelism** requires `tensor_model_parallel_size > 1` +- **Expert parallelism** requires MoE models (`num_moe_experts > 0`) +- All parallelism strategies can be combined, but total parallelism must divide evenly into the world size + +## Resources + +- [Megatron Core Developer Guide](https://docs.nvidia.com/megatron-core/developer-guide/latest/) +- [Scaling Language Model Training](https://developer.nvidia.com/blog/scaling-language-model-training-to-a-trillion-parameters-using-megatron/) +- [Megatron-LM Repository](https://github.com/NVIDIA/Megatron-LM) +- [Transformer Engine](https://github.com/NVIDIA/TransformerEngine) diff --git a/fern/v0.2.0/pages/performance-guide.mdx b/fern/v0.2.0/pages/performance-guide.mdx new file mode 100644 index 0000000000..dc365522ce --- /dev/null +++ b/fern/v0.2.0/pages/performance-guide.mdx @@ -0,0 +1,708 @@ +--- +title: "Performance Tuning Guide" +description: "" +--- + +Megatron-Bridge provides a wide range of features for performant and memory-efficient LLM training on GPUs, and comes pre-configured with optimal settings. However, factors such as model architecture, hyperparameters, GPU count, and GPU type can affect the available options, and additional tuning may be necessary to achieve optimal performance. This document explores the factors that affect training performance, highlights common issues, and outlines techniques for performance tuning that lead to higher MFU (Model FLOPS Utilization) and TCO. + + +This guide makes references to several configuration settings. These settings will be referenced relative to the the config class that contains them, e.g. `OptimizerConfig.lr`. Please see [API Documentation](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/) for more details on configuration settings. + + + +This guide references several configuration settings from `TransformerConfig`. Please apply these to the appropriate ModelProvider for your model, e.g. `GPTModelProvider`, as the `ConfigContainer` does not accept a raw `TransformerConfig`. + + +## Low Precision Training + +1. Expected speedup of FP8 training compared to BF16 training + + > 1. The default low-precision LLM training recipe applies FP8 computation exclusively to the linear layers within the Transformer block, typically achieving a speedup of 1.2–1.5X. + > 2. However, the actual speedup depends on the proportion of training time spent on these linear layers. For instance, smaller LLMs with a limited hidden size exhibit lower FP8 speedup, as linear layers scale with O(sequence_length × hidden_size²) complexity, whereas the other element-wise computation layers (e.g., layer norms, dropouts, RoPE, and simple math functions) scale with O(sequence_length × hidden_size), and dot-product attention scales with O(sequence_length² × hidden_size). Consequently, the contribution of linear layers to the overall training time is smaller in such models. + > 3. Different FP8 recipes use varying quantization block sizes, affecting performance. Smaller quantization blocks generally incur higher overhead in both quantization and GEMM execution. For example, MXFP8 with a 1×32 quantization block performs less efficiently than full tensor-wise FP8 scaling. + +2. Common issues of low FP8 training speedup + + > 1. Host performance boundness when LLM uses small GPU kernels (see [Lowering Host Overhead and Jitters](#lowering-overhead-jitter)). + > 2. A low proportion of linear layers in training step time that use FP8 computation. + +## Parallel Mapping Strategies + +1. Data Parallelism using Distributed Optimizer + + > 1. You should begin with data-parallel (DP) mapping. As long as the model and activation memory fit within the GPUs, data parallelism generally offers optimal performance, minimizes communication overhead, and maximizes per-GPU tensor sizes (compared to per-tensor sharding). + > + > 2. Megatron-Bridge uses the distributed optimizer as the default method for data-parallel training. It shards master parameters and optimizer states across data-parallel ranks, reducing model state memory usage without increasing communication overhead compared to traditional data-parallel training. + > + > > 1. `OptimizerConfig.use_distributed_optimizer=true` + +2. Per-tensor Sharding (Tensor-parallel or Context-parallel mappings) + + > 1. Tensor parallelism (TP) is the primary recommendation when a model exceeds GPU memory capacity under data-parallel mapping. However, since it involves higher communication overhead, the tensor-parallel size should ideally be confined to the high-bandwidth intra-node network (NVLink domain). + > + > > 1. `TransformerConfig.tensor_model_parallel_size=` + > + > 2. When the sequence length in a training run is significantly larger than the hidden size, activation memory can overflow. In such cases, context parallelism (CP) helps by sharding tensors along the sequence dimension, allowing the workload to fit within limited GPU memory and improving performance. Like tensor parallelism (TP), CP requires inter-GPU communication of activations. However, for the same tensor sizes, CP generally results in lower communication volume. + +That said, CP’s effectiveness depends on the relative sizes of the sequence length and hidden size. When the sequence length is smaller than the hidden size, CP produces narrow (or "skinny") tensor shards on each GPU. This reduces data reuse and can degrade performance. + +Additionally, because CP shards activations, it also partitions optimizer states in distributed training. As a result, optimizer state partitioning spans both the data parallel (DP) and context parallel (CP) dimensions. + +> > 1. `TransformerConfig.context_parallel_size=` +> +> 1. Performance tips: +> +> > 1. A large tensor-parallel or context-parallel size is not recommended unless the hidden size or sequence length is large enough to maintain sufficient per-GPU parallelism and avoid excessive communication overhead. For example, using a tensor-parallel size of 8 for LLAMA 3 70B could lead to low GPU utilization and make training host-performance bound. +> > 2. You can combine TP and CP to optimize performance by balancing communication overhead. For example, using TP=2 along with CP=2 can give better performance than TP=4 when the sequence size is larger than the hidden size. +> > 3. For additional tips, see [Long Sequence Training](#long-sequence-train). + +1. Pipeline Parallelism + + > 1. Pipeline parallelism (PP) is necessary when a model cannot fit within GPU memory using tensor parallelism. Also, virtual pipeline parallelism (VPP) should be used in conjunction with pipeline parallelism to reduce the overhead caused by pipeline warm-up and flush bubbles. + > + > > 1. `TransformerConfig.pipeline_model_parallel_size=` + > > 2. `TransformerConfig.virtual_pipeline_model_parallel_size=` + > + > 2. Performance tips in PP and VPP sizing: + > + > > 1. PP can also be combined with per-tensor sharding methods to mitigate the impact of sharding inefficiencies and pipeline bubbles. For instance, TP4 + PP2 may outperform TP8 when both mappings fit into memory because using a large TP reduces per-GPU tensor sizes but increases the communication cost, increasing the exposed communication. + > > 2. VPP increases inter-stage communication overhead. When a global batch contains many micro-batches, using a smaller VPP size can improve performance, as the exposed communication cost outweighs the reduction in pipeline bubbles. + > + > 3. Asymmetric Transformer layer allocation across pipeline stages + > + > > 1. An LLM with a large vocabulary size has computationally heavy embedding lookup and projection operations, leading to load imbalance across pipeline stages. To address this, Megatron-Bridge provides an option to allocate one fewer Transformer layer in the first and last pipeline stages, which handle embedding lookup and projection, to better balance workloads. + > > + > > > 1. `GPTProvider.account_for_embedding_in_pipeline_split=true` + > > > 2. `GPTProvider.account_for_loss_in_pipeline_split=true` + +2. Expert Parallelism + + > 1. Expert Parallelism (EP) is designed specifically for Mixture-of-Experts (MoE) models to efficiently distribute sparse MLP weights across multiple chips. It can be used in combination with other parallelism strategies such as Tensor Parallelism (TP), Context Parallelism (CP), Pipeline Parallelism (PP), Data Parallelism (DP), and Fully Sharded Data Parallel (FSDP). In the current design, the dense attention part and the sparse MLP part are fully decoupled in terms of their TP, CP, and DP parallelism configurations. Expert Tensor Parallelism (ETP) is introduced to specifically control the tensor parallelism for the sparse MLP part. ETP uses TP for dense layers for the ranks allocated for EP in sparse layers. On the other hand, the baseline is DEP, which folds DP in dense layers for EP in sparse layers. + > + > > 1. `TransformerConfig.expert_model_parallel_size=` + > > 2. `TransformerConfig.expert_tensor_parallel_size=` + > + > 2. Performance tips in hybrid folding options and EP sizing: + > + > > 1. Typically, EP is kept within the high-bandwidth intra-node network (NVLink domain) to minimize the communication overhead it can introduce. However, using communication overlap techniques—such as pipeline overlap or 1F1B overlap—along with PP (e.g., DualPipe) might make it possible to expand EP into the inter-node networks. + > > + > > 2. Within the sparse MLP block, DP replaces CP because it has no impact on the computation pattern based on the dispatched tokens in each EP rank. + > > + > > 3. Usually, ETP is set to 1 to avoid significant communication overhead that comes with applying TP to MLP GEMMs. + > > + > > 4. When multiple experts are placed on a single chip after applying Expert Parallelism, enabling grouped GEMM can significantly improve computation efficiency. + > > + > > > 1. `TransformerConfig.moe_grouped_gemm=True` + +3. Fully Sharded Data Parallelism + + > 1. Megatron-Bridge supports PyTorch-native FSDP. FSDP can be used in combination with per-tensor sharding methods. + > + > > 1. To use PyTorch FSDP2: + > > + > > > 1. `DistributedInitConfig.use_torch_fsdp2=True` + > + > 2. FSDP can be preferred over TP+PP+DP mappings in the following scenarios: + > + > > 1. Small models with a large sequence, thus the parameter AllGather and gradient ReduceScatter can effectively be hidden under computation and the short communication overlap causes minor interference to the computation under overlap. + > > 2. In FSDP training, activation storage remains as the main memory bottleneck because FSDP only shards model state memory, and a large per-GPU activation is needed to hide the costly FSDP communication. On GB200 GPUs, Megatron-Bridge offers an option to offload activations to the host memory via a high-speed chip-to-chip interconnect. + > > 3. Baseline training is host performance-bound, but FSDP allows for larger per-GPU tensor sizes by eliminating TP or enabling a larger micro-batch size. + + {/* TODO: support megatron custom fsdp */} + {/* > 1. Megatron-Bridge supports two Fully Sharded Data Parallelism (FSDP) implementations: PyTorch-native FSDP and a custom Megatron FSDP built within Megatron Core. While both follow the same sharding principles, the custom implementation is further optimized for performance. The performance gain of the custom FSDP comes primarily from minimizing the data movement to the communication tensors and reusing communication buffers. Both FSDP methods can be used in combination with per-tensor sharding methods. */} + {/* > */} + {/* > > 1. To use PyTorch FSDP2: */} + {/* > > */} + {/* > > > 1. `DistributedInitConfig.use_torch_fsdp2=True` */} + {/* > > */} + {/* > > 2. To use Custom Megatron FSDP: */} + {/* > > */} + {/* > > > 1. `recipe.trainer.strategy.fsdp="megatron"` */} + {/* > > > 2. `recipe.trainer.strategy.ddp.data_parallel_sharding_strategy="optim_grads_params"` */} + {/* > */} + {/* > 2. FSDP can be preferred over TP+PP+DP mappings in the following scenarios: */} + {/* > */} + {/* > > 1. Small models with a large sequence, thus the parameter AllGather and gradient ReduceScatter can effectively be hidden under computation and the short communication overlap causes minor interference to the computation under overlap. */} + {/* > > 2. In FSDP training, activation storage remains as the main memory bottleneck because FSDP only shards model state memory, and a large per-GPU activation is needed to hide the costly FSDP communication. On GB200 GPUs, Megatron-Bridge offers an option to offload activations to the host memory via a high-speed chip-to-chip interconnect. */} + {/* > > 3. Baseline training is host performance-bound, but FSDP allows for larger per-GPU tensor sizes by eliminating TP or enabling a larger micro-batch size. */} + +4. Heterogeneous Encoder Parallelism + + > 1. Encoder Pipeline Parallel + > + > > 1. Use `T5ModelProvider.encoder_pipeline_model_parallel_size`. + > > 2. In an Encoder-Decoder architecture like Multimodal models (VLMs like NeVA etc.), Encoder Pipeline Parallel can be used to add pipeline parallelism to the encoder. + > > 3. Pipeline parallelism controls the amount of pipelining in the decoder part. + > > 4. Encoder Pipeline Parallel is limited to 1 at the moment, i.e., the encoder can occupy a maximum of 1 PP stage. + > > 5. By default, Encoder Pipeline Parallel is 0 and Decoder Pipeline Parallel is 1. + > > 6. When the Encoder Pipeline Parallel size is 0, it shares the first PP stage of the Decoder. + > + > 2. Encoder Tensor Parallel + > + > > 1. Use `T5ModelProvider.encoder_tensor_model_parallel_size`. + > > 2. Since encoders tend to be much smaller than decoders, we also provide the ability to set a different amount of tensor parallelism to the encoder than the decoder. + > > 3. By default, encoder tensor parallel is set to 0, i.e., the amount of tensor parallelism in the encoder is equal to tensor parallelism in the decoder. + > > 4. To use this option, Encoder Pipeline Parallel must be greater than 0 as we need the encoder to be on its own pipeline stage. + > > 5. Encoder Tensor Parallel size is limited to be less than or equal to Tensor parallel size. + > + > 3. Total number of GPUs required when these features are used is: + > + > > 1. Data Parallel size * Context Parallel size * ((Encoder TP * Encoder PP) + (Decoder TP * Decoder PP)) + > + > 4. These features are experimental and may still have bugs. There are critical bug fixes that will be made in a future release. + +5. Parallel mapping strategies with NVL72 + + > 1. Training with only data parallelism or FSDP makes it straightforward to fully utilize the bandwidth of an NVL72 system. However, when combining multiple parallelism strategies, it's important to ensure that high-volume communicators remain confined within each NVL72 domain. For example, with TP=4, DP=16, and PP=4, the GPUs in the first TP group of DP1/PP1 spans both NVLink and network domains, causing communication performance to be bottlenecked by the slower network link. To avoid this, you may choose TP and DP sizes such that the product of TP × DP divides evenly into the NVL72 configuration. If the model-parallel size does not align naturally, padding may be required to support non-divisible group sizes. + > 2. To avoid this partitioning complexity, you can just use 64 GPUs out of the 72 GPUs. + +## Communication Overlaps and Tuning + +1. Data-parallel communication of Distributed Optimizer + + > 1. Distributed optimizer overlaps parameter AllGathers with the forward computation of the first micro-batch and gradient ReduceScatters with the backward computation of the last micro-batch. + > + > > 1. `DistributedDataParallelConfig.overlap_param_gather=true` + > > 2. `DistributedDataParallelConfig.overlap_grad_reduce=true` + > + > 2. When using the distributed optimizer with pipeline parallelism (PP) + virtual pipeline parallelism (VPP), DP communications overlap with multiple micro-batches, increasing the opportunity for effective overlap. Also, Megatron-Bridge aligns the execution timing of DP communications across pipeline-parallel ranks to synchronize the computing kernel slowdown from the overlap. + > + > > 1. `DistributedDataParallelConfig.align_param_gather=true` + > + > 3. Slow DP communication at large scaling training: + > + > > 1. Distributing optimizer states across a partial DP domain reduces communication costs over high-latency Ethernet networks. Model states remain replicated outside the distributed domain. During the final micro-batch backpropagation, gradient ReduceScatters occur within the distributed domain, followed by AllReduce in the non-distributed domain. Parameter AllGathers are performed only within the distributed domain. + > > + > > > 1. `DistributedDataParallelConfig.num_distributed_optimizer_instances= ` + > > + > > 2. A large message size for DP communication is recommended to maximize network bandwidth utilization. You can achieve this by increasing the communication bucket size. + > > + > > > 1. `DistributedDataParallelConfig.bucket_size=` + > + > 4. A common reason for DP communication overlap failure: + > + > > 1. Persistent Layer Normalization (LN) kernels from Transformer Engine use spin-waiting for all SMs in the GPU, causing the LN kernel and subsequent computation kernels to be scheduled only after DP communication. To prevent this, an appropriate SM margin should be configured using the following environment variables. + > > + > > > 1. `NVTE_FWD_LAYERNORM_SM_MARGIN=<#SM for DP collectives = 16>` + > > > 2. `NVTE_BWD_LAYERNORM_SM_MARGIN=<#SM for DP collectives = 16>` + +{/* 2. Custom Megatron FSDP */} + +{/* > 1. Unless you specify the communication bucket size, MCORE FSDP uses fixed communication overlap that overlaps the parameter AllGather and gradient ReduceScatter of each Transformer layer with its associated forward and backward computations. */} + +3. Tensor-parallel (TP) communication (with sequence parallelism) + + > 1. Megatron-Bridge currently uses the userbuffer backend in Transformer Engine for TP communication overlaps. This offers the pipelined overlap of the TP communication with dependent computation. + > + > > 1. `CommOverlapConfig.tp_comm_overlap` + > + > 2. The overlap method, resource, and precision of the TP communication overlaps are configurable, and the most performant configurations are set in the Megatron-Bridge training recipes by default. Also, you can set a custom TP communication overlap configuration via the below interface following the structure of TransformerLayerTPOverlapCfg class. + > + > > 1. `CommOverlapConfig.tp_comm_overlap_cfg=` + > + > 3. TP communication overlap setting tips + > + > > 1. Balancing the number of SMs between communication and GEMM + > > + > > > 1. For AllGather/ReduceScatter bulk and ReduceScatter pipelined overlap, you can adjust the number of SMs to balance communication and GEMM execution. Allocating too many SMs to communication may degrade GEMM performance, while too few may expose communication overhead. The default SM allocation for communication is 16, but you can fine-tune it based on profiling results. + > > > 2. `TPOverlapCfg.num_sm=` + > > + > > 2. CGA sizing to improve SM utilization + > > + > > > 1. The CGA size can be set between 1 and 4, but it should not exceed the number of SMs allocated for communication. We recommend using CGA ≤ 2 to prevent potential SM rasterization that could impact GEMM performance. + > > > 2. `TPOverlapCfg.cga_size=` + > > + > > 3. Use 4× splits for ReduceScatter and GEMM overlap to optimize the balance between GEMM efficiency and communication exposure. + > > + > > > 1. In GEMM-then-ReduceScatter pipeline overlap, a 1× ReduceScatter chunk remains exposed. A small split size increases communication exposure, while a large split size may degrade performance due to aggregated GEMM wave quantization. We find that num_splits = 4 generally provides the best performance. + > > > 2. `TPOverlapCfg.num_split=` + > + > 4. Common reason for TP comm overlap failure at Hopper + > + > > 1. At H100 GPU, an environment variable `CUDA_DEVICE_MAX_CONNECTIONS=1` should be set. Otherwise, TP communication kernels can be scheduled at the end of GEMM to overlap with. + > > 2. Pipelined TP communication overlap is used by a static userbuffer registered upon model initialization. Therefore, it doesn't support activation tensors dynamically changing between steps or between Transformer layers. + +4. Context-parallel (CP) communication + + > 1. CP communication is configurable via "cp_comm_type", which can be "p2p", "all_gather", "a2a", or "a2a+p2p". Communications of "p2p" are implemented as ring-exchange send/receive operations, and they are hard-coded to overlap with the attention compute of sequence chunks. See [Long Sequence Training](#long-sequence-train) for more details. + +5. Expert-parallel communication + + > 1. To hide the A2A/AG communication introduced by EP, pipeline split overlap or 1F1B overlap alongside Pipeline Parallelism could be possible. It will be added to Megatron-Bridge in future releases. + +6. Pipeline-parallel (PP) send/receive communication + + > 1. PP send/recv in steady 1F1B states are set to be overlapped with computes by default. + > 2. The PP send/recv in warmup and flush are exposed by default. + +(comm-data-types)= +## Communication Data Types + +1. FP8 data-parallel parameter AllGather in Distributed Optimizer and FSDP + + > 1. Megatron-Bridge supports FP8 parameter AllGather for per-tensor FP8 scaling recipes. This operation is lossless, enhancing performance while reducing memory usage. + > + > > 1. `MixedPrecisionConfig.fp8_param=true` + +2. BF16 (instead of FP32) data-parallel reduction in Distributed Optimizer and FSDP + + > 1. We have validated that BF16 reduction is numerically safe across numerous model training runs. However, BF16 reduction with a large data-parallel size (e.g., DP ≥ 128), especially the Ring reduction algorithm—which accumulates copies sequentially—may impact numerical stability. When using SHARP with NVIDIA InfiniBand, BF16 reduction is more robust, as it performs binary additions with higher precision for intermediate partial reductions. + > + > > 1. `DistributedDataParallelConfig.grad_reduce_in_fp32=false` + +3. FP8 tensor-parallel ReduceScatter + + > 1. When communication latency exceeds GEMM execution time, using FP8 input ReduceScatter can better hide communication overhead. This approach has low numerical impact, as the GEMM output must be cast to FP8 and then converted back to high precision during reduction. + > + > > 1. `TPOverlapCfg.fp8_buf=true` + +4. FP8 A2A Dispatch for expert parallel communication + + > 1. Megatron-Bridge is working on supporting FP8 A2A dispatch (before expert FC1), but still keeps BF16 A2A combine (after expert FC2). + +## Performance at Scale + +1. Scaling a training job is typically achieved by increasing the size of the data-parallel domain. In large-scale training, this often results in a small number of micro-batches per global batch—or even a single micro-batch—causing most computations to overlap with data-parallel communication. To maintain high performance in such scenarios, you should focus on minimizing the overhead of data-parallel communication and reducing host-driven inter-GPU jitter. + +2. You can lower the overhead of data-parallel communication by (1) reducing the communication precision e.g., BF16 for gradient reduction and FP8 parameter gathering, (2) improving the efficiency of communication by increasing the data-parallel communication message size or using the hierarchical data-parallel reduction, or (3) using multi-cast and switch reduction with SHARP in case of InfiniBand network. + + > 1. Using BF16 gradient reduction and FP8 parameter gather are described in [Communication Data Types](#comm-data-types) + > + > 2. For non-pipeline-parallel training, the data-parallel communication bucket size can be adjusted using the knobs below. In pipeline-parallel training, however, the bucket size is fixed and determined by the number of parameters assigned to each virtual pipeline rank. + > + > > 1. `DistributedDataParallelConfig.bucket_size=` + > + > 3. Setting the knob below splits the data-parallel domain of the distributed optimizer into a sharding domain and a replication domain. Gradient reduction then occurs in two stages—one within each domain—avoiding the use of a single large flat ring for collective operations that have high latency. + > + > > 1. `DistributedDataParallelConfig.num_distributed_optimizer_instances=` + +3. Ideas to reduce the host-driven inter-GPU jitters are discussed in [Lowering Host Overhead and Jitters](#lowering-overhead-jitter). + +(lowering-overhead-jitter)= +## Lowering Host Overhead and Jitters + +1. Common observation associated with host overhead + + > 1. Significantly low GPU FLOPS. + > 2. Small performance gain of low-precision (FP8) training. + > 3. Small LLMs with small hidden size or sequence length or fine-tuning without sequence packing + > 4. High multi-GPU communication variation. + +2. Increasing micro-batch size and reduce per-tensor sharding + + > 1. The most common way to increase per-GPU tensor size is by increasing the micro-batch size or minimizing unnecessary per-tensor sharding (e.g., TP or CP) when GPU memory permits. + +3. Manual garbage collection to align the host interruption across GPUs + + > 1. Megatron-Bridge manually aligns the timing of garbage collection across GPUs that significantly mitigate the host overhead compared to the baseline automatic garbage collection. + > + > > 1. `TrainingConfig.manual_gc_interval=` + +4. CUDA graph to eliminate repeated static host code execution + + > 1. Megatron-Bridge supports graph capture, significantly reducing host overhead. CUDA Graph is applicable only to LLMs with a static tensor shape across training steps. For example, it supports fixed-size packed sequences but does not handle sequences with varying lengths at each step. Also, MoE models with token-dropless propagation have limited CUDA graph support, restricted to the dense modules only. + > 2. CUDA graph requires additional memory for static buffer management, typically adding a few gigabytes for static buffers, while models with PP size > 1 may consume over 10GB. We are actively working to reduce this memory overhead. + > 3. `TransformerConfig.enable_cuda_graph=true` + +5. Bind CPU memory for GPU processes + + > 1. Binding CPU cores to GPU processes helps mitigate long latency issues and ensures minimal variation in GPU queuing latency across GPUs. This optimization significantly impacts, particularly when the communication domain size is large. + > 2. Example command line for a X86-based GPU system: `numactl --cpunodebind=$((SLURM_LOCALID/4)) --membind=$((SLURM_LOCALID/4)) ` + > 3. Example command line for a Grace-based GPU system: `numactl --cpunodebind=$((SLURM_LOCALID/2)) --membind=$((SLURM_LOCALID/2)) ` + +(reducing-memory-overflow)= +## Techniques for Reducing Memory to Avoid Memory Overflow and Enhance Training Efficiency + +1. Activation recomputation + + > 1. Megatron-Bridge LLMs default to dot-product attention-only recomputation using Flash Attention, efficiently regenerating large intermediate activations from the attention operation with minimal computational overhead. + > + > 2. Megatron-Bridge also supports recomputing the full intermediate activations of a Transformer block, significantly reducing activation memory usage at the cost of approximately 30% additional computation. The number of Transformer blocks to recompute can be adjusted using a configurable setting. + > + > > 1. `TransformerConfig.recompute_granuality=full` + > > 2. `TransformerConfig.recompute_method=block` + > > 3. `TransformerConfig.recompute_num_layers=` + +2. Activation offloading to host memory + + > 1. Megatron-Bridge supports offloading activation memory to host memory, essential for training tasks constrained by activation memory. This is particularly useful for scenarios like (1) FSDP, where model state memory is minimized through sharding but activation memory remains high, (2) LoRA, which has frozen parameters but significant activation memory demands, and (3) the training with a large sequence length. The efficiency of activation offloading depends on both the interconnect bandwidth between the GPU and host and the host memory bandwidth. From this perspective, Grace-based systems like the GB200 enhance offloading performance by optimizing these bandwidths. + > + > 2. The following knobs should be configured to enable offloading and specify the number of Transformer layers to offload to host memory. The maximum number of layers that can be offloaded depends on host memory capacity, which may be lower when the CPU is shared among multiple GPUs. + > + > > 1. `TransformerConfig.cpu_offloading=True` + > > 2. `TransformerConfig.cpu_offloading_weights=False` + > > 3. `TransformerConfig.cpu_offloading_num_layers= ` + > + > 3. Environment variable settings to avoid resource conflict between CPU memory offloading and network communication + > + > > 1. `NCCL_NET_GDR_LEVEL=PHB # NCCL <=2.25` + > > 2. `NCCL_NET_GDR_C2C=1 # NCCL >=2.26` + > + > 4. Optimization tips + > + > > 1. Given the ratio between activation volume and computational operations, offloading all layer activations naively can become a performance bottleneck. Optimizing performance requires tuning the number of layers to offload while balancing it with recomputation. + +3. Weight memory-optimized BF16 training + + > 1. In BF16 training, Megatron-Bridge optimizes memory usage by storing only the BF16 remainder of the master weight copies for the next optimizer update. This is possible because BF16 data can be represented using a subset of FP32 bits, allowing Megatron-Bridge to avoid redundant storage of the FP32 portion used for BF16 representation. This is default enabled when using precision-aware optimizer in Megatron Core. + > + > > 1. `OptimizerConfig.use_precision_aware_optimizer=True` + +4. Common memory usage hikes from environment variable setting + + > 1. The below environment variables will (1) avoid preserving the buffers for NCCL communication and (2) disable NVLSharp when not used. Both these options lower the GPU memory usage. + > + > > 1. `TORCH_NCCL_AVOID_RECORD_STREAMS=1` + > > 2. `NCCL_NVLS_ENABLE=0` + > + > 2. While not enabled by default, you can further reduce memory usage caused by segmentation penalties by setting the env var shown below. + > + > > 1. `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` + +5. Keep parameters in FP8 at FP8 training + + > 1. In FP8 training, after optimizer step execution, we can keep the parameters in FP8. Compared to the baseline that keeps the intermediate weight values in BF16, FP8 parameters lower memory usage and improve communication performance. The below knob enables keeping the parameters in FP8. + > + > > 1. `MixedPrecisionConfig.fp8_param_gather=True` + +## Operator Fusion + +1. You can control specific fusion behaviors using the following configuration knobs: + + > 1. `TransformerConfig.masked_softmax_fusion=true` + > 2. `GPTProvider.cross_entropy_loss_fusion=true` + > 3. `GPTProvider.gradient_accumulation_fusion=true` + > 4. `TransformerConfig.bias_activation_fusion=true` + > 5. `TransformerConfig.bias_dropout_fusion=true` + > 6. `TransformerConfig.apply_rope_fusion=true` + +2. Megatron-Bridge offers different Flash Attention options, which can be chosen through the model config: + + > 1. Let Transformer Engine decide (default): `TransformerConfig.attention_backend=AttnBackend.auto` + > 2. FlashAttention2: `TransformerConfig.attention_backend=AttnBackend.flash` + > 3. cuDNN fused attention: `TransformerConfig.attention_backend=AttnBackend.fused` + +(long-sequence-train)= +## Long Sequence Training + +1. Problem of long sequence training + + > 1. Training with long sequence length can lead to memory overflow due to the huge memory cost of activations. The problem could be solved by recomputing activations in backward, but it can impose up to ~30% overheads in each training step. Context parallelism is a better solution which splits the sequence dimension across multiple GPUs, so that each GPU only computes and saves activations of a sequence chunk. In this way, memory overflow is addressed without introducing any redundant compute. + +2. CP to shard activation (knob) + + > 1. `TransformerConfig.context_parallel_size=` + > + > > 1. Both TP and CP can reduce activation memory overheads. It's not wise to be biased to either of them. Communications of TP and CP are overlapped by GEMM and Attention respectively. Blindly enlarging their sizes can make some communications hard to overlap. It's recommended to sweep a combination of TP+CP configs. The optimal config is expected to make full use of all related compute and do best overlapping, thereby achieving best end-to-end performance. + > + > 2. `TransformerConfig.cp_comm_type= or ` + > + > > 1. Megatron-Core provides multiple implementation variants of CP and allows you to make choices based on your specific use cases by configuring "cp_comm_type". The configuration value can be `p2p`, `all_gather`, `a2a`, or `a2a+p2p`. These communication types are compatible with each other, so they can be flexibly interleaved between transformer layers. You only need to provide a list, where each element corresponds to a layer. + > > 2. `p2p`: exchanges KV sequence chunks in ring-topology. The P2P communications can be fully overlapped. + > > 3. `all_gather`: inserts an all-gather before attention to get a full sequence of KV. The all-gather is exposed, but it should not impose big overheads if GQA/MQA are used, as they have very few KV heads. + > > 4. `a2a`: is an implementation of DeepSpeed Ulysses. A2A communications are added before and after the attention module to gather full sequence length and further scatter heads in CP domain. A2A cannot be overlapped. + > > 5. `a2a+p2p`: is a middle ground between `a2a` and `p2p`. This is useful for cases of big CP sizes, where each sequence chunk is too short to overlap P2P communications. It first does A2A in partial CP groups to gather relatively longer sequence chunks, then applies P2P implementation to the gathered chunks. It also can be helpful for hierarchical CP communications, for example A2A and P2P happen in NVLink and IBLink domains respectively. + > > 6. With small and medium CP size, `p2p` is the recommended configuration because communications can be fully overlapped; "all_gather" also should work fine with GQA/MQA. As for strongly-scaling a sequence length with big CP sizes, the short chunk length can barely overlap the `p2p` communications, so `a2a+p2p` ought to be the preferred choice. `a2a` could be adopted in some cases for its simplicity. However, CP size can be restricted with "a2a" because it requires the number of attention heads to be divisible by CP size. Restricted CP size will finally limit the sequence length that can be run. + +3. Activation recomputation (in [Techniques for Reducing Memory to Avoid Memory Overflow and Enhance Training Efficiency](#reducing-memory-overflow)) + +4. Activation offloading to host memory (in [Techniques for Reducing Memory to Avoid Memory Overflow and Enhance Training Efficiency](#reducing-memory-overflow)) + +## Sequence Packing for Performant Fine-Tuning + +1. Dataset preparation + + > 1. Fine-tuning datasets with shorter sequences of variable length can be packed into longer sequences, up to a set maximum length, for best efficiency. + +2. To use this feature, the microbatch size must be set to 1. In place of increasing the micro batch size, the maximum sequence length can be increased, which will effectively increase the number of individual sequences per packed sequence. + +3. Enabled with: + + > 1. `FinetuningDatasetConfig.packed_sequence_specs.packed_sequence_size=` + > 2. `TrainingConfig.micro_batch_size=1` + +4. Performance benefits also include: + + > 1. Inconsistent lengths between sequences in the fine-tuning dataset would reduce the computation efficiency. With a micro-batch size over 1, all sequences must be padded with empty tokens to the length of the longest one in the micro-batch. Similarly, some optimizations like CUDA graphs require uniform sequence lengths between micro-batches. Packed sequences are arranged so that the total number of tokens per packed sequence is as close to the maximum length as possible, making most processed tokens useful. + > 2. Likewise, when using data parallel, variance in time needed to process different batches can result in all batches needing to wait for the longest to finish-- and this variance is reduced with packed sequence. + +## GPU Core Clock Optimization + +1. Increase the clock ratio of GPU core over off-chip memory system + + > 1. NVIDIA GPUs support a CPU core clock boost mode, which increases the core clock rate by reducing the off-chip memory clock rate. This is particularly beneficial for LLMs, which are typically compute throughput-bound. + > + > > 1. `sudo nvidia-smi boost-slider --vboost 1 ` + +## Profiling Options for Analysis-based Performance Tuning + +1. Nsight system profile + + > 1. Megatron-Bridge provides an interface to enable the NVIDIA Nsight Systems profiler, which displays the GPU execution trace of all CUDA streams. You can check whether communication kernels overlap with computation kernels and adjust resource allocation to balance communication and computation. The Nsight Systems profile can be enabled using ProfilingConfig, as shown below. + > 2. `ProfilingConfig(use_nsys_profiler=True, profile_start_step=, profile_end_step=, profile_ranks=<[0,...]>)` + +2. Memory snapshot + + > 1. Megatron-Bridge provides an interface to extract the memory snapshot that shows the memory allocation bytes, the allocation lifespan, and the function call stack. Extracting the memory snapshot can be enabled by ProfilingConfig as shown below. + > 2. `ProfilingConfig(record_memory_history=True, memory_snapshot_path=)` + +## DeepEP: Common Issues and Solutions + +DeepEP is a communication library optimized for Mixture-of-Experts (MoE) all-to-all operations. When using DeepEP for cross-node Expert Parallelism (EP), there are several common issues related to network transport and GPU-NIC affinity that can significantly impact performance. + +> Note: DeepEP is best optimized for NVL8 systems such as the DGX-B200 NVL8 or DGX-H200 NVL8. For GB200 NVL72 rack-scale systems, where 72 GPUs are interconnected within the same NVLINK domain, we recommend using [HybridEP](https://github.com/deepseek-ai/DeepEP/tree/hybrid-ep) instead of DeepEP. HybridEP is maintained by NVIDIA and is specifically optimized for NVL72 rack scale systems. It is also integrated into the Megatron-core [fused all-to-all module](https://docs.nvidia.com/megatron-core/developer-guide/latest/apidocs/core/core.transformer.moe.fused_a2a.html) as an alternative backend under the `flex` token dispatcher. +> +> Learn more about GB200 MoE training best practices [here](https://github.com/NVIDIA/Megatron-LM/blob/dev/docs/discussions/deepseek-v3-gb200-optimization/deepseek-v3-gb200-reproduce-guide.md). + +### 1. Why is my DeepEP not working + +1. What is IBGDA and why is it a problem + + DeepEP achieves optimal cross-node communication performance using InfiniBand GPU Direct Async (IBGDA), which is supported by ConnectX NICs in both InfiniBand and RoCEv2 modes. However, IBGDA is not always enabled by default—it often requires cluster administrators to actively configure the system and enable GPU Direct RDMA support in the InfiniBand (or RoCEv2) fabric. If this configuration step is skipped or unsupported in the cluster environment, IBGDA may be unavailable, which can prevent DeepEP inter-node EP capability from functioning. + +1. Network Transport: IBGDA vs. IBRC + + > 1. IBGDA (InfiniBand GPU Direct Async) requires cluster administrators to enable GPU Direct RDMA and configure the InfiniBand subsystem. Many clusters do not have IBGDA enabled by default. + > 2. The official DeepEP main branch has removed support for IBRC (InfiniBand Reliable Connection), which previously served as a fallback mechanism. With IBRC, a CPU proxy thread will assist in processing the EP communication, which might have performance degradation compared to IBGDA, but we find such performance degradation doesn't overshadow the benefit of enabling wideEP in production training. + +2. Solution: NVSHMEM 3.5 with Automatic Transport Fallback + + > 1. NVSHMEM 3.5 introduces improved auto-fallback support for cross-node communication under various network configurations. It can automatically select the best available transport (IBGDA, IBRC, or other supported mechanisms) based on cluster capabilities. + > 2. To benefit from NVSHMEM’s auto-fallback in DeepEP: + > - Download the [official NVSHMEM 3.5.19-1 release](https://github.com/NVIDIA/nvshmem/releases/tag/v3.5.19-1). You can also choose to compile it from source in your container environment; we provide such examples later in this guide. + > - Switch to the [DeepEP branch with native NVSHMEM API integration](https://github.com/seth-howell/DeepEP/tree/nvshmem_native_apis). This branch enables automatic use of NVSHMEM’s fallback mechanisms without requiring any manual code modifications. + +### 2. GPU-NIC Affinity and Bandwidth Contention + +A common cause of poor DeepEP performance is incorrect GPU-to-NIC (Network Interface Card) affinity, where multiple GPUs compete for bandwidth on a single NIC. As noted in [DeepEP PR #466](https://github.com/deepseek-ai/DeepEP/pull/466), cross-node EP performance may degrade if multiple GPUs use the same NIC, due to certain GPU-NIC affinity in some clusters. This PR provides a solution by supporting the environment variable `DEEP_EP_DEVICE_TO_HCA_MAPPING` to specify GPU-to-NIC mappings so that each GPU is automatically bound to the optimal NIC for maximum DeepEP throughput. + +With this PR's solution, we need the following environment variables to map GPUs to NICs correctly. First, you need to find out the names of the NICs by running `ibstat`. In our example, we found the following for one RoCEv2 DGX-B200 cluster: +``` +> ibstat | grep ^CA +CA 'rocep145s0' +CA 'rocep146s0' +CA 'rocep152s0' +CA 'rocep153s0' +CA 'rocep198s0' +CA 'rocep199s0' +CA 'rocep205s0' +CA 'rocep206s0' +``` + +Use the following environment variables to map GPUs to NICs. Note that `0:rocep145s0:1` is formatted as `::` so that each GPU will only be mapped to one dedicated NIC. +```bash +export NVSHMEM_ENABLE_NIC_PE_MAPPING=1 +export DEEP_EP_DEVICE_TO_HCA_MAPPING="0:rocep145s0:1,1:rocep146s0:1,2:rocep152s0:1,3:rocep153s0:1,4:rocep198s0:1,5:rocep199s0:1,6:rocep205s0:1,7:rocep206s0:1" +``` + +### 3. Build DeepEP + +In this section, we provide a reference Dockerfile that shows how to build NVSHMEM 3.5 and the customized DeepEP into your container environment. + +Note that the following example is provided for DGX-B200 NVL8 systems, but similar ideas apply to Hopper generation as well—just change the Dockerfile accordingly. For example, you just need to change the compile target for SM90. + +Key points: + +- NVSHMEM source: https://github.com/NVIDIA/nvshmem/tree/v3.5.19-1 +- DeepEP branch that we cherry-picked with all the fixes above: https://github.com/zhongbozhu/DeepEP/tree/nvshmem_deepep_gcp +- Example training container template for DGX-B200: https://github.com/yanring/Megatron-MoE-ModelZoo/blob/main/dockers/B200.Dockerfile + +**Dockerfile** +```bash +FROM nvcr.io/nvidia/pytorch:25.11-py3 as base + +# Other dependencie you may want +... + +# Dependency of IBGDA +RUN ln -s /usr/lib/x86_64-linux-gnu/libmlx5.so.1 /usr/lib/x86_64-linux-gnu/libmlx5.so + +# Clone DeepEP customized version +WORKDIR /home/dpsk_a2a +RUN git clone https://github.com/zhongbozhu/DeepEP.git ./deepep +RUN cd ./deepep && git checkout nvshmem_deepep_gcp && cd /home/dpsk_a2a + +# Clone NVSHMEM 3.5 https://github.com/NVIDIA/nvshmem +RUN git clone --branch v3.5.19-1 https://github.com/NVIDIA/nvshmem.git ./deepep-nvshmem +RUN cd ./deepep-nvshmem && git checkout v3.5.19-1 && cd /home/dpsk_a2a + +# Build nvshmem from source +# You can also download the pre-built binary, and skip the following +RUN apt-get update && \ + DEBIAN_FRONTEND=noninteractive apt-get install -y \ + clang \ + llvm-dev \ + libclang-dev && \ + rm -rf /var/lib/apt/lists/* + +WORKDIR /home/dpsk_a2a/deepep-nvshmem +RUN mkdir -p build && mkdir -p install && \ + cmake -S . -B build \ + -DCMAKE_INSTALL_PREFIX=/home/dpsk_a2a/deepep-nvshmem/install \ + -DCUDA_HOME=/usr/local/cuda \ + -DMPI_HOME=/opt/hpcx/ompi \ + -DMPI_C_COMPILER=/opt/hpcx/ompi/bin/mpicc \ + -DMPI_CXX_COMPILER=/opt/hpcx/ompi/bin/mpicxx \ + -DNVSHMEM_MPI_SUPPORT=OFF \ + -DNVSHMEM_IBRC_SUPPORT=ON \ + -DNVSHMEM_IBGDA_SUPPORT=ON \ + -DNVSHMEM_IBDEVX_SUPPORT=OFF \ + -DNVSHMEM_UCX_SUPPORT=OFF \ + -DNVSHMEM_SHMEM_SUPPORT=OFF \ + -DNVSHMEM_PMIX_SUPPORT=OFF \ + -DNVSHMEM_USE_NCCL=OFF \ + -DNVSHMEM_USE_GDRCOPY=ON \ + -DGDRCOPY_HOME=/usr \ + -DNVSHMEM_USE_MLX5DV=ON \ + -DNVSHMEM_BUILD_TESTS=ON \ + -DNVSHMEM_BUILD_EXAMPLES=ON \ + -DNVSHMEM_BUILD_PYTHON_LIB=OFF \ + -DNVSHMEM_BUILD_BITCODE_LIBRARY=OFF \ + -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_CUDA_ARCHITECTURES="100" && \ + cmake --build build -j && \ + cmake --install build + +ENV NVSHMEM_DIR=/home/dpsk_a2a/deepep-nvshmem/install +ENV LD_LIBRARY_PATH=${NVSHMEM_DIR}/lib:$LD_LIBRARY_PATH +ENV PATH=${NVSHMEM_DIR}/bin:$PATH + +## Build deepep +WORKDIR /home/dpsk_a2a/deepep +ENV TORCH_CUDA_ARCH_LIST="10.0" +ENV PIP_NO_BUILD_ISOLATION=1 +ENV CPATH=${CUDA_HOME}/include/cccl:$CPATH +RUN pip install --no-build-isolation . + +``` + +DeepEP provides `test_internode.py` to test and benchmark cross-node EP communication. In our experiment, when using 4 nodes of DGX-B200 (i.e., EP32), the achieved throughput for cross-EP is about 50 GB/s with IBRC. We provide an example SLURM script below for running such a test with DeepEP. + +In another experiment on the same cluster, with IBGDA enabled by the cluster admin, we observed approximately 10% higher inter-node performance—roughly 55 GB/s. To enable IBGDA, you need to set the environment variable `export NVSHMEM_IB_ENABLE_IBGDA=true`; there is no need to change the software version or container, because with the software provided above, both modes will work. + +```bash +srun --account= -N 4 -p batch --time 30 \ + --ntasks-per-node=1 --gpus-per-node=8 \ + --no-container-mount-home --container-mounts "/lustre:/lustre" \ + --container-image \ + --mpi=none --export=ALL \ + bash -lc ' +set -eo pipefail + +# Env Var for GPU-NIC mapping +export NVSHMEM_ENABLE_NIC_PE_MAPPING=1 +export DEEP_EP_DEVICE_TO_HCA_MAPPING="0:rocep145s0:1,1:rocep146s0:1,2:rocep152s0:1,3:rocep153s0:1,4:rocep198s0:1,5:rocep199s0:1,6:rocep205s0:1,7:rocep206s0:1" + +# 1) Expand SLURM_JOB_NODELIST and grab the first hostname +headnode=$(python - </dev/null 2>&1; then + master_ip=$(getent ahostsv4 "$headnode" | awk "{print \$1; exit}") +else + master_ip="" +fi +MASTER_ADDR="${master_ip:-$headnode}" + +# 3) Export rendezvous env that matches test_internode.py expectations +export MASTER_ADDR +export MASTER_PORT=${MASTER_PORT:-29500} +export WORLD_SIZE=${SLURM_NNODES:-2} # number of nodes +export RANK=${SLURM_NODEID:-0} # 0..N-1 per node + +export OMP_NUM_THREADS=1 +python -u /home/dpsk_a2a/deepep/tests/test_internode.py +' + +``` + +## Index - List of Tuning Knobs + +- `CommOverlapConfig.tp_comm_overlap` +- `CommOverlapConfig.tp_comm_overlap_cfg` +- `CUDA_DEVICE_MAX_CONNECTIONS` +- `TrainingConfig.manual_gc_interval` +- `MixedPrecisionConfig.fp8_param` +- `ProfilingConfig` +- `NCCL_NET_GDR_C2C` +- `NCCL_NET_GDR_LEVEL` +- `NCCL_NVLS_ENABLE` +- `NVTE_BWD_LAYERNORM_SM_MARGIN=<#SM for DP collectives` +- `TransformerConfig.attention_backend` +- `AttnBackend` +- `NVTE_FWD_LAYERNORM_SM_MARGIN=<#SM for DP collectives` +- `PYTORCH_CUDA_ALLOC_CONF` +- `TrainingConfig.micro_batch_size` +- `FinetuningDatasetConfig.packed_sequence_specs.packed_sequence_size` +- `TransformerConfig.apply_rope_fusion` +- `TransformerConfig.bias_activation_fusion` +- `TransformerConfig.bias_dropout_fusion` +- `TransformerConfig.cp_comm_type` +- `TransformerConfig.cpu_offloading` +- `TransformerConfig.cpu_offloading_num_layers` +- `TransformerConfig.cpu_offloading_weights` +- `GPTProvider.cross_entropy_loss_fusion` +- `TransformerConfig.enable_cuda_graph` +- `MixedPrecisionConfig.fp8_param_gather` +- `GPTProvider.gradient_accumulation_fusion` +- `TransformerConfig.masked_softmax_fusion` +- `TransformerConfig.recompute_granuality` +- `TransformerConfig.recompute_method` +- `TransformerConfig.recompute_num_layers` +- `OptimizerConfig.use_precision_aware_optimizer` +- `GPTProvider.account_for_embedding_in_pipeline_split` +- `GPTProvider.account_for_loss_in_pipeline_split` +- `TransformerConfig.context_parallel_size` +- `DistributedDataParallelConfig.align_param_gather` +- `DistributedDataParallelConfig.bucket_size` +- `DistributedDataParallelConfig.bucket_size` +- `DistributedDataParallelConfig.data_parallel_sharding_strategy` +- `DistributedDataParallelConfig.grad_reduce_in_fp32` +- `DistributedDataParallelConfig.num_distributed_optimizer_instances` +- `DistributedDataParallelConfig.overlap_grad_reduce` +- `DistributedDataParallelConfig.overlap_param_gather` +- `T5ModelProvider.encoder_pipeline_model_parallel_size` +- `T5ModelProvider.encoder_tensor_model_parallel_size` +- `TransformerConfig.expert_model_parallel_size=` +- `TransformerConfig.expert_tensor_parallel_size=` +- `TransformerConfig.moe_grouped_gemm` +- `DistributedInitConfig.use_torch_fsdp2` +- `TransformerConfig.pipeline_model_parallel_size` +- `TransformerConfig.tensor_model_parallel_size` +- `TransformerConfig.virtual_pipeline_model_parallel_size` +- `OptimizerConfig.use_distributed_optimizer` +- `TORCH_NCCL_AVOID_RECORD_STREAMS` +- `TPOverlapCfg.cga_size` +- `TPOverlapCfg.fp8_buf` +- `TPOverlapCfg.num_sm` +- `TPOverlapCfg.num_split` +{/* - `garbageCollectionCallback.gc_interval_val` */} +{/* - `NsysPlugin` */} diff --git a/fern/v0.2.0/pages/performance-summary.mdx b/fern/v0.2.0/pages/performance-summary.mdx new file mode 100644 index 0000000000..f69af0f092 --- /dev/null +++ b/fern/v0.2.0/pages/performance-summary.mdx @@ -0,0 +1,135 @@ +--- +title: "Performance" +description: "" +--- + +As part of the NVIDIA NeMo Framework, Megatron Bridge, provides optimal performance for training advanced generative AI models by incorporating the most recent training techniques, such as model parallelization, optimized attention mechanisms, and more, to achieve high training throughput. + +This page provides performance benchmarks for large language models using Megatron-Bridge across different GPU systems and configurations. + +## Nomenclature + +- **GBS**: Global Batch Size +- **MBS**: Micro Batch Size +- **FSDP**: Fully Sharded Data Parallel + - FSDP = 1: use FSDP + - FSDP = 0: use DDP (Distributed Data Parallel) +- **TP**: Tensor Parallel Size +- **PP**: Pipeline Parallel Size +- **CP**: Context Parallel Size +- **VP**: Virtual Pipeline Parallel Size +- **EP**: Expert Parallel Size +- **GA**: Number of Gradient Accumulations + +## Performance Metrics + +Performance is measured using: + +- **Tokens/sec/GPU**: Throughput per GPU +- **Model TFLOP/sec/GPU**: Model floating-point operations per second per GPU + +## Performance Summary for Large Language Models + +Below are performance benchmarks for various large language models organized by release version. These results were obtained using performance recipes available [here](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/scripts/performance). + +The performance data includes: + +- **Pre-training Performance**: Throughput metrics for various model sizes and architectures +- **System Configurations**: Results across different GPU systems (DGX-GB200, DGX-B200, DGX-H100) +- **Precision Options**: Performance comparisons between different precision modes (BF16, FP8, MXFP8) + +--- + +## 25.11 NeMo Container + +### Pre-Training Performance + +#### System: DGX-GB300 + +| Model | #-GPUs | Precision | GBS | MBS | Sequence Length | FSDP | TP | PP | CP | VP | EP | GA | Tokens / sec / GPU | Model TFLOP / sec / GPU | +|-------|--------|-----------|-----|-----|-----------------|------|----|----|----|----|----|----|-----------------------|-------------------------| +| LLAMA3_8B | 8 | FP8-CS (FP8-MX) | 128 | 2 | 8192 | 0 | 1 | 1 | 1 | n/a | 1 | 8 | 37556 (36108) | 1933 (1858) | +| LLAMA3_70B | 64 | FP8-CS | 256 | 2 | 8192 | 1 | 1 | 1 | 1 | 1 | 1 | 2 | 4520 | 2030 | +| LLAMA3.1_405B | 256 | FP8-CS | 1536 | 1 | 8192 | 0 | 2 | 8 | 2 | 4 | 1 | 192 | 999 | 2522 | +| DeepSeekV3 (w/o MTP) | 256 | BF16 | 4096 | 1 | 4096 | 0 | 1 | 2 | 1 | 4 | 32 | 32 | 3848 | 961 | +| DeepSeekV3 (w/o MTP)| 256 | FP8-MX | 4096 | 1 | 4096 | 0 | 1 | 2 | 1 | 4 | 32 | 32 | 4357 | 1088 | +| GPT OSS 120B | 64 | BF16 | 1280 | 2 | 8192 | 0 | 1 | 1 | 1 | 1 | 64 | 10 | 18347 | 565 | +| Qwen3_30B_a3B | 8 | FP8-MX | 512 | 4 | 4096 | 0 | 1 | 1 | 1 | 1 | 8 | 16 | 28934 | 666 | +| Qwen3_235B_a22B | 256 | BF16 | 8192 | 2 | 4096 | 0 | 1 | 4 | 1 | 12 | 16 | 32 | 6131 | 907 | + +#### System: DGX-GB200 + +| Model | #-GPUs | Precision | GBS | MBS | Sequence Length | FSDP | TP | PP | CP | VP | EP | GA | Tokens / sec / GPU | Model TFLOP / sec / GPU | +|-------|--------|-----------|-----|-----|-----------------|------|----|----|----|----|----|----|-----------------------|-------------------------| +| LLAMA3_8B | 8 | FP8-CS (FP8-MX) | 128 | 2 | 8192 | 0 | 1 | 1 | 1 | n/a | 1 | 8 | 31508 (29789) | 1622 (1533) | +| LLAMA3_70B | 64 | FP8-CS | 256 | 2 | 8192 | 1 | 1 | 1 | 1 | 1 | 1 | 2 | 4312 | 1937 | +| LLAMA3.1_405B | 256 | FP8-CS | 1536 | 1 | 8192 | 0 | 4 | 16 | 1 | 4 | 1 | 384 | 813 | 2053 | +| DeepSeekV3 (w/o MTP) | 256 | BF16 | 4096 | 1 | 4096 | 0 | 1 | 4 | 1 | 4 | 64 | 64 | 3139 | 782 | +| DeepSeekV3 (w/o MTP) | 256 | FP8-MX | 4096 | 1 | 4096 | 0 | 1 | 8 | 1 | 4 | 32 | 128 | 4018 | 1003 | +| GPT OSS 120B | 64 | BF16 | 1280 | 1 | 8192 | 0 | 1 | 1 | 1 | 1 | 64 | 20 | 15876 | 488 | +| Qwen3_30B_a3B | 8 | FP8-MX | 512 | 4 | 4096 | 0 | 1 | 1 | 1 | 1 | 8 | 16 | 23766 | 547 | +| Qwen3_235B_a22B | 256 | BF16 | 8192 | 1 | 4096 | 0 | 1 | 8 | 1 | 3 | 32 | 256 | 4916 | 728 | + +#### System: DGX-B200 + +| Model | #-GPUs | Precision | GBS | MBS | Sequence Length | FSDP | TP | PP | CP | VP | EP | GA | Tokens / sec / GPU | Model TFLOP / sec / GPU | +|-------|--------|-----------|-----|-----|-----------------|------|----|----|----|----|----|----|-----------------------|-------------------------| +| LLAMA3_8B | 8 | FP8-CS (FP8-MX) | 128 | 2 | 8192 | 0 | 1 | 1 | 1 | n/a | 1 | 8 | 30624 (29521) | 1576 (1519) | +| LLAMA3.1_405B | 128 | FP8-CS (FP8-MX) | 64 | 1 | 8192 | 0 | 4 | 8 | 2 | 8 | 1 | 32 | 661 (624) | 1667 (1576) | +| DeepSeekV3 (w/ MTP) | 256 | FP8-MX | 2048 | 1 | 4096 | 0 | 1 | 16 | 1 | 1 | 8 | 128 | 2139 | 557 | +| GPT OSS 120B | 64 | BF16 | 512 | 4 | 4096 | 0 | 1 | 1 | 1 | 1 | 8 | 2 | 8213 | 223 | +| Qwen3_30B_a3B | 8 | FP8-MX | 512 | 1 | 4096 | 0 | 1 | 1 | 1 | 1 | 8 | 64 | 9299 | 214 | +| Qwen3_235B_a22B | 64 | FP8-MX | 1024 | 1 | 4096 | 0 | 1 | 8 | 1 | 2 | 8 | 128 | 3269 | 484 | + +#### System: DGX-H100 + +| Model | #-GPUs | Precision | GBS | MBS | Sequence Length | FSDP | TP | PP | CP | VP | EP | GA | Tokens / sec / GPU | Model TFLOP / sec / GPU | +|-------|--------|-----------|-----|-----|-----------------|------|----|----|----|----|----|----|-----------------------|-------------------------| +| LLAMA3_8B | 8 | FP8-CS | 128 | 1 | 8192 | 1 | 1 | 1 | 1 | n/a | 1 | 16 | 14451 | 744 | +| LLAMA3_70B | 64 | FP8-CS | 128 | 1 | 8192 | 0 | 4 | 8 | 1 | 5 | 1 | 64 | 1602 | 719 | +| LLAMA3.1_405B | 1024 | FP8-CS | 512 | 1 | 8192 | 0 | 8 | 8 | 2 | 8 | 1 | 64 | 292 | 737 | +| GPT OSS 120B | 64 | BF16 | 512 | 4 | 4096 | 0 | 1 | 4 | 1 | 1 | 8 | 2 | 5630 | 153 | +| Qwen3_30B_a3B | 16 | FP8-CS | 512 | 2 | 4096 | 0 | 1 | 2 | 1 | 24 | 8 | 32 | 5275 | 121 | +| Qwen3_235B_a22B | 256 | FP8-CS | 2048 | 1 | 4096 | 0 | 2 | 8 | 1 | 4 | 32 | 128 | 1575 | 233 | + +- The numbers in normal parentheses indicate the use of different quantization granularities: In case of GB200 and B200 systems, 32×32 for both weights and activations. For H100 system, 128×128 for weights and 1×128 for activations, which match those used in the original DeepSeekV3 pre-training. +- In MoE training benchmarks, we force-balance the token distribution among experts and all benchmarks are token-dropless. + +## 25.09 NeMo Container + +### Pre-Training Performance + +#### System: DGX-GB200 + +| Model | #-GPUs | GBS | MBS | Sequence Length | FSDP | TP | PP | CP | VP | EP | GA | Tokens / sec / GPU | Model TFLOP / sec / GPU | +|-------|--------|-----|-----|-----------------|------|----|----|----|----|----|----|-----------------------|-------------------------| +| LLAMA3_8B | 8 | 128 | 2 | 8192 | 0 | 1 | 1 | 1 | n/a | 1 | 8 | 31357 (29925) | 1614 (1540) | +| LLAMA3_70B | 64 | 128 | 2 | 8192 | 1 (0) | 1 (2) | 1 (4) | 1 | 1 (5) | 1 | 1 (16) | 3986 (3546) | 1791 (1593) | +| LLAMA3.1_405B | 128 | 64 | 1 | 8192 | 1 (0) | 2 (4) | 1 (8) | 1 (2) | 1 (8) | 1 | 1 (32) | 729 (578) | 1840 (1458) | +| DeepSeekV3 (tokendrop) | 256 | 2048 | 1 | 4096 | 0 | 1 | 4 (8) | 1 | 4 (2) | 64 | 32 (64) | 3454 (2835) | 899 (738) | +| Qwen3_30B_a3B (tokendrop) | 8 | 512 | 4 | 4096 | 0 | 1 | 1 | 1 | 1 | 8 | 16 | 22775 (23723) | 524 (546) | +| Qwen3_235B_a22B (tokendrop) | 64 | 1024 | 1 | 4096 | 0 | 2 | 1 | 1 | 1 | 64 | 32 | 4452 (4416) | 659 (654) | + +#### System: DGX-B200 + +| Model | #-GPUs | GBS | MBS | Sequence Length | FSDP | TP | PP | CP | VP | EP | GA | Tokens / sec / GPU | Model TFLOP / sec / GPU | +|-------|--------|-----|-----|-----------------|------|----|----|----|----|----|----|-----------------------|-------------------------| +| LLAMA3_8B | 8 | 128 | 2 | 8192 | 0 | 1 | 1 | 1 | n/a | 1 | 8 | 29994 (29388) | 1544 (1513) | +| LLAMA3.1_405B | 128 | 64 | 1 | 8192 | 0 | 4 | 8 | 2 | 8 | 1 | 32 | 664 (622) | 1676 (1569) | +| DeepSeekV3 (tokendrop) | 256 | 2048 | 1 | 4096 | 0 | 1 | 16 | 1 | 1 | 8 | 128 | 2265 (2159) | 589 (562) | +| Qwen3_30B_a3B (tokendrop) | 8 | 512 | 1 | 4096 | 0 | 1 | 1 | 1 | 1 | 8 | 64 | 18066 | 416 | +| Qwen3_235B_a22B (tokendrop) | 64 | 1024 | 1 | 4096 | 0 | 1 | 8 | 1 | 2 | 8 | 128 | 4104 (4275) | 607 (633) | + +#### System: DGX-H100 + +| Model | #-GPUs | GBS | MBS | Sequence Length | FSDP | TP | PP | CP | VP | EP | GA | Tokens / sec / GPU | Model TFLOP / sec / GPU | +|-------|--------|-----|-----|-----------------|------|----|----|----|----|----|----|-----------------------|-------------------------| +| LLAMA3_8B | 8 | 128 | 1 | 8192 | 1 | 1 | 1 | 1 | n/a | 1 | 16 | 14079 | 725 | +| LLAMA3_70B | 64 | 128 | 1 | 8192 | 0 | 4 | 8 | 1 | 5 | 1 | 64 | 1619 | 727 | +| LLAMA3.1_405B | 1024 | 512 | 1 | 8192 | 0 | 8 | 8 | 2 | 8 | 1 | 64 | 302 | 763 | +| DeepSeekV3 (dropless) | 1024 | 8192 | 1 | 4096 | 0 | 2 | 8 | 1 | 4 | 64 | 128 | 1297 | 338 (330) | +| Qwen3_30B_a3B (tokendrop) | 16 | 512 | 2 | 4096 | 0 | 1 | 2 | 1 | 24 | 8 | 32 | 10494 | 241 | +| Qwen3_235B_a22B (tokendrop) | 256 | 2048 | 1 | 4096 | 0 | 2 | 8 | 1 | 4 | 32 | 128 | 1204 | 178 | + +- The numbers in parentheses indicate the use of different quantization granularities: In case of GB200 and B200 systems, 32×32 for both weights and activations. For H100 system, 128×128 for weights and 1×128 for activations, which match those used in the original DeepSeekV3 pre-training. +- In token-dropless MoE training benchmarks, we force-balance the token distribution among experts. diff --git a/fern/v0.2.0/pages/recipe-usage.mdx b/fern/v0.2.0/pages/recipe-usage.mdx new file mode 100644 index 0000000000..f86003a97f --- /dev/null +++ b/fern/v0.2.0/pages/recipe-usage.mdx @@ -0,0 +1,263 @@ +--- +title: "Using Recipes" +description: "" +--- + +Megatron Bridge provides production-ready training recipes for several popular models. You can find an overview of supported recipes and 🤗 HuggingFace bridges [here](/index#supported-models). +This guide will cover the next steps to make use of a training recipe, including how to [override configuration](#overriding-configuration) and how to [launch a job](#launch-methods). + +## Overview + +- **Coverage**: We provide recipes across select model families and sizes, including Llama, Qwen, DeepSeek, and Nemotron-H (Mamba-based). +- **Defaults**: Each recipe sets defaults meant for convergence and performance across parallelisms, precision data types, and optimizer & scheduler choices. These recipes can be used as a high-quality starting point. +- **Integration**: Recipes return a single `ConfigContainer` that plugs directly into our training [entry points](/training/entry-points) (see the published docs as well: https://docs.nvidia.com/nemo/megatron-bridge/latest/training/entry-points.html). +- **Customization**: You can override any part of the recipe (Python, YAML, CLI) to adapt to your data, scale, and objectives. + +## Overriding configuration + +Recipes are provided through a [`bridge.training.config.ConfigContainer`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.ConfigContainer.html) object. This is a dataclass that holds all configuration objects needed for training. You can find a more detailed overview of the `ConfigContainer` [here](/training/config-container-overview). +The benefit of providing the full recipe through a pythonic structure is that it is agnostic to any configuration approach that a user may prefer, whether that's YAML, `argparse` or something else. In other words, the user may override the recipe however they see fit. + +The following sections detail a few different ways to override the configuration recipe. For a complete training script, please see [this example](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/models/llama/pretrain_llama3_8b.py). + +### Python + +If you prefer to manage configuration in Python, you can directly modify attributes of the `ConfigContainer`: + +```python +from megatron.bridge.recipes.llama.llama3_8b import pretrain_config + +# Get the base ConfigContainer from the recipe +cfg: ConfigContainer = pretrain_config() + +# Apply overrides. Note the hierarchical structure +cfg.train.train_iters = 20 +cfg.train.global_batch_size = 8 +cfg.train.micro_batch_size = 1 +cfg.logger.log_interval = 1 +``` + +You can also replace entire sub-configs of the `ConfigContainer`: + +```python +from megatron.bridge.recipes.llama.llama3_8b import pretrain_config +from megatron.bridge.models.llama import Llama3ModelProvider + +cfg: ConfigContainer = pretrain_config() + +small_llama = Llama3ModelProvider( + num_layers=2, + hidden_size=768, + ffn_hidden_size=2688, + num_attention_heads=16, +) +cfg.model = small_llama +``` + +### YAML +Overriding a configuration recipe with a YAML file can be done using OmegaConf utilities: + +```python +from omegaconf import OmegaConf +from megatron.bridge.recipes.llama.llama3_8b import pretrain_config +from megatron.bridge.training.utils.omegaconf_utils import ( + apply_overrides, + create_omegaconf_dict_config, +) + +cfg: ConfigContainer = pretrain_config() +yaml_filepath = "conf/llama3-8b-benchmark-cfg.yaml" + +# Convert the initial Python dataclass to an OmegaConf DictConfig for merging +# excluded_fields holds some configuration that cannot be serialized into a DictConfig +merged_omega_conf, excluded_fields = create_omegaconf_dict_config(cfg) + +# Load and merge YAML overrides +yaml_overrides_omega = OmegaConf.load(yaml_filepath) +merged_omega_conf = OmegaConf.merge(merged_omega_conf, yaml_overrides_omega) + +# Apply overrides while preserving excluded fields +final_overrides_as_dict = OmegaConf.to_container(merged_omega_conf, resolve=True) +apply_overrides(cfg, final_overrides_as_dict, excluded_fields) +``` + +The above snippet will update `cfg` with all overrides from `llama3-8b-benchmark-cfg.yaml`. + +### Hydra-style + +Megatron Bridge provides some utilities to update the ConfigContainer using Hydra-style CLI overrides: + +```python +import sys +from omegaconf import OmegaConf +from megatron.bridge.recipes.llama.llama3_8b import pretrain_config +from megatron.bridge.training.utils.omegaconf_utils import ( + apply_overrides, + create_omegaconf_dict_config, + parse_hydra_overrides, +) + +cfg: ConfigContainer = pretrain_config() +cli_overrides = sys.argv[1:] + +# Convert the initial Python dataclass to an OmegaConf DictConfig for merging +# excluded_fields holds some configuration that cannot be serialized into a DictConfig +merged_omega_conf, excluded_fields = create_omegaconf_dict_config(cfg) + +# Parse and merge CLI overrides +merged_omega_conf = parse_hydra_overrides(merged_omega_conf, cli_overrides) + +# Apply overrides while preserving excluded fields +final_overrides_as_dict = OmegaConf.to_container(merged_omega_conf, resolve=True) +apply_overrides(cfg, final_overrides_as_dict, excluded_fields) +``` + +After the above snippet, `cfg` will be updated with all CLI-provided overrides. +A script containing the above code could be called like so: + +```sh +torchrun pretrain_cli_overrides.py model.tensor_model_parallel_size=4 train.train_iters=100000 ... +``` + +## Launch methods + +Megatron Bridge supports launching scripts with both `torchrun` and [NeMo-Run](https://github.com/NVIDIA-NeMo/Run). +Once your script is ready to be launched, refer to one of the following sections. + +### Torchrun +Megatron Bridge training scripts can be launched with the `torchrun` command that most PyTorch users are familiar with. +Simply specify the number of GPUs to use with `--nproc-per-node` and the number of nodes with `--nnodes`. For example, on a single node: + +```sh +torchrun --nnodes 1 --nproc-per-node 8 /path/to/train/script.py +``` + +For multi-node training, it is recommended to use a cluster orchestration system like SLURM. +The `torchrun` command should be wrapped as specified by your cluster orchestration system. +For example, with Slurm, wrap the `torchrun` command inside of `srun`: + +```sh +# launch.sub + +srun --nodes 2 --gpus-per-node 8 \ + --container-image --container-mounts \ + bash -c " + torchrun --nnodes $SLURM_NNODES --nproc-per-node $SLURM_GPUS_PER_NODE /path/to/train/script.py + " +``` + +Along with any other required flags. It is also recommended to use a NeMo Framework container with Slurm. You can find a list of container tags on [NGC](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/nemo/tags). + +### NeMo-Run + +Megatron Bridge also supports launching training with [NeMo-Run](https://docs.nvidia.com/nemo-framework/user-guide/latest/nemorun/index.html). NeMo-Run is a Python package that enables configuring and executing experiments across several platforms. +For multi-node training, NeMo-Run will generate a script with appropriate commands, similar to the `srun` command described above. + +The recommended method to launch a Megatron Bridge script with NeMo-Run is through the `run.Script` API. +You can modify the following 3 steps to your needs in a new file: + +```python +import nemo_run as run + +if __name__ == "__main__": + # 1) Configure the `run.Script` object + train_script = run.Script(path="/path/to/train/script.py", entrypoint="python") + + # 2) Define an executor for the desired target platform + executor = run.LocalExecutor(ntasks_per_node=8, launcher="torchrun") + + # 3) Execute + run.run(train_script, executor=executor) +``` + +NeMo-Run supports launching on several different platforms, including [SLURM clusters](https://docs.nvidia.com/nemo-framework/user-guide/latest/nemorun/guides/execution.html#slurmexecutor). +For more details, please see the NeMo-Run [documentation](https://docs.nvidia.com/nemo-framework/user-guide/latest/nemorun/guides/execution.html#) for a list of supported platforms, their corresponding executors, and configuration instructions. + +You can also forward arguments from the NeMo-Run launch script to the target script: + +```python +import nemo_run as run +import argparse + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + ... + known_args, args_to_fwd = parser.parse_known_args() + train_script = run.Script(..., args=args_to_fwd) +``` + +For a complete example of the `run.Script` API, including argument forwarding, please see [this script](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/models/llama/pretrain_llama3_8b_nemo_run_script.py). + +#### Plugins + +Megatron Bridge provides several NeMo-Run plugins to simplify the usage of certain features. +These plugins can simply be added to the `run.run()` call: + +```python +import nemo_run as run +from megatron.bridge.recipes.run_plugins import NsysPlugin + +if __name__ == "__main__": + train_script = run.Script(path="/path/to/train/script.py", entrypoint="python") + executor = run.LocalExecutor(ntasks_per_node=8, launcher="torchrun") + + plugins = [] # plugins argument expects a list + nsys = NsysPlugin(profile_step_start=10, profile_step_end=15, ...) + plugins.append(nsys) + run.run(train_script, plugins=plugins, executor=executor) +``` + +##### Custom Argument Converters + +By default, plugins convert their configuration to Hydra-style CLI arguments when used with `run.Script` tasks. If your training script uses a different argument format (e.g., argparse), you can provide a custom converter function via the `script_args_converter_fn` parameter. + +```python +import nemo_run as run +from typing import List +from megatron.bridge.recipes.run_plugins import ( + PreemptionPlugin, + PreemptionPluginScriptArgs, +) + +# Define a custom converter for argparse-style arguments +def argparse_preemption_converter(args: PreemptionPluginScriptArgs) -> List[str]: + result = [] + if args.enable_exit_handler: + result.append("--enable-exit-handler") + if args.enable_exit_handler_for_data_loader: + result.append("--enable-exit-handler-dataloader") + return result + +if __name__ == "__main__": + train_script = run.Script(path="/path/to/train/script.py", entrypoint="python") + executor = run.LocalExecutor(ntasks_per_node=8, launcher="torchrun") + + # Use the plugin with the custom converter + plugin = PreemptionPlugin( + preempt_time=120, + enable_exit_handler=True, + script_args_converter_fn=argparse_preemption_converter, + ) + run.run(train_script, plugins=[plugin], executor=executor) +``` + +Each plugin provides its own corresponding dataclass (e.g., `PreemptionPluginScriptArgs`, `NsysPluginScriptArgs`) that defines the available arguments for conversion. + +See the [API reference](#bridge.recipes.run_plugins) for a list of available NeMo-Run plugins. + +### Avoiding Hangs + +When working with any scripts in Megatron Bridge, please make sure you wrap your code in an `if __name__ == "__main__":` +block. Otherwise, your code may hang unexpectedly. + +The reason for this is that Megatron Bridge uses Python's `multiprocessing` module in the backend when running a +multi-GPU job. The multiprocessing module will create new Python processes that will import the current module (your +script). If you did not add `__name__== "__main__"`, then your module will spawn new processes which import the +module and then each spawn new processes. This results in an infinite loop of process spawning. + +## Resources + +- [OmegaConf documentation](https://omegaconf.readthedocs.io/en/2.3_branch/) +- [torchrun Documentation](https://docs.pytorch.org/docs/stable/elastic/run.html) +- [PyTorch Multinode Training documentation](https://docs.pytorch.org/tutorials/intermediate/ddp_series_multinode.html) +- [NeMo-Run documentation](https://docs.nvidia.com/nemo-framework/user-guide/latest/nemorun/index.html#) diff --git a/fern/v0.2.0/pages/releases/README.mdx b/fern/v0.2.0/pages/releases/README.mdx new file mode 100644 index 0000000000..84fb2285df --- /dev/null +++ b/fern/v0.2.0/pages/releases/README.mdx @@ -0,0 +1,81 @@ +--- +title: "Releases" +description: "" +--- + +This directory contains release information, version history, and known issues for Megatron Bridge. + +## Release Documentation + +| Document | Purpose | When to Read | +|----------|---------|--------------| +| **[Software Versions](/software-versions)** | Current software versions and dependencies | Checking compatibility, planning upgrades | +| **[Changelog](/changelog)** | Detailed release history and changes | Understanding what changed in each release | +| **[Known Issues](/known-issues)** | Known bugs, limitations, and workarounds | Troubleshooting issues, planning workarounds | + +## Quick Navigation + +### I want to + +**📋 Check current versions** +→ See [Software Versions](/software-versions) for current versions and dependencies + +**📝 Review release history** +→ Read [Changelog](/changelog) for detailed release notes + +**🐛 Find known issues** +→ Check [Known Issues](/known-issues) for bugs and workarounds + +**🔄 Plan an upgrade** +→ Review [Changelog](/changelog) and [Known Issues](/known-issues) before upgrading + +**🔍 Troubleshoot a problem** +→ Check [Known Issues](/known-issues) for reported issues and solutions + +## Release Information Overview + +### Software Versions + +The [Software Versions](/software-versions) document provides: + +- Current Megatron Bridge version +- Required dependency versions +- Compatible software versions +- Version compatibility matrix + +### Changelog + +The [Changelog](/changelog) includes: + +- Release dates and version numbers +- New features and enhancements +- Bug fixes and improvements +- Breaking changes and migration notes +- Deprecations and removals + +### Known Issues + +The [Known Issues](/known-issues) document covers: + +- Known bugs and their status +- Limitations and constraints +- Workarounds and solutions +- Planned fixes and timelines + +## Related Documentation + +- **[Main Documentation](/../README)** - Return to main documentation +- **[Documentation Guide](/../documentation)** - Contributing to documentation + +## Staying Up to Date + +To stay informed about releases: + +1. **Check Software Versions** - Review [Software Versions](/software-versions) regularly +2. **Read Changelog** - Review [Changelog](/changelog) for each release +3. **Monitor Known Issues** - Check [Known Issues](/known-issues) for updates +4. **Follow Release Notes** - Review release announcements for important changes + +--- + +**Need version information?** Check [Software Versions](/software-versions) or return to the [main documentation](/../README). diff --git a/fern/v0.2.0/pages/releases/changelog.mdx b/fern/v0.2.0/pages/releases/changelog.mdx new file mode 100644 index 0000000000..c24d81084d --- /dev/null +++ b/fern/v0.2.0/pages/releases/changelog.mdx @@ -0,0 +1,119 @@ +--- +title: "Changelog" +description: "" +--- + +## NVIDIA Megatron-Bridge 0.2.2 + +* This release addresses known security issues. For the latest NVIDIA Vulnerability Disclosure Information visit [https://www.nvidia.com/en-us/security/](https://www.nvidia.com/en-us/security/), for acknowledgement please reach out to the NVIDIA PSIRT team at [PSIRT@nvidia.com](mailto:PSIRT@nvidia.com) + +## NVIDIA Megatron-Bridge 0.2.1 + +* Performance + * Activation offloading to host memory support with pipelining + * Supports the high activation memory needs of MoE models training with dynamic shapes + * Fixed Nemotron FLOPS calculation model +* Model Collection Support + * Ministral 3 +* Enhanced LoRA support + * LoRA support for Mamba layers (for Nemotron Nano V2 and NemotronH finetuning) + +## NVIDIA Megatron-Bridge 0.2.0 + +* [Model Collection Support](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models) + + * LLM + * HuggingFace Conversion + training recipes: + * GPT-OSS + * Qwen3 Next + * Nemotron-H + * Nemotron Nano v2 + * Moonlight + * OlMoE + * GLM 4.5 + * Gemma 3 + * HuggingFace conversion support: + * Llama Nemotron + * Mistral + * Gemma + * Gemma 2 + * VLM + * Nemotron Nano v2 VL + * Qwen 3 VL + * Qwen2.5 VL + * Gemma3 VL + +* [Performance](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/scripts/performance) + * Megatron-Bridge support for new benchmarks + * Benchmarks (same workloads as GB200 system) for GB300 system + * GPT-OSS 120B + * Qwen3-Next 80B_A3B + * Support for linear attention on Blackwell - Gated Delta Networks + * Pre-training with NVFP4 precision: Llama3 8B, Lama3 70B, Llama3.1 405B + * Megatron-Bridge support for benchmarks previously existing only for NeMo 2.0 + * Nemotron-H 56B + * Fine-tuning (SFT and LoRA): Llama3 8B and Llama3 70B + * HybridEP: DeepSeek V3 benchmarks on GB200 and GB300 systems now use HybridEP + * CUDA Graphs + * Full-model iteration CUDA graph used for dense models- Llama3 8B, Llama3 70B, Llama3.1 405B + * Fine-grained Transformer component specific CUDA Graphs used for MoE models + +* [NVIDIA Model Optimization Integration](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/examples/quantization) + * Knowledge Distillation + * Post training quantization export + * Quantization aware training + +* [Enhanced LoRA support](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/peft) + * Support for expert layers + * Supported merging adapters for export to HuggingFace @HollowMan6 + +* [Finetuning dataset improvements: OpenAI messages format conversion, chat template support](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/data) +* [Integration with Tensor NVIDIA-DLFW-Inspect for tensor statistic collection & monitoring](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/training/tensor_inspect.py) +* [Support for sample-based training](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/training/config.py) +* Broader Community Adoption: Integrate the Megatron-Bridge into the training pipelines of VeRL ([PR](https://github.com/volcengine/verl/pull/4063/files)), Slime ([PR](https://github.com/THUDM/slime/pull/894/)), and Sky-RL ([PR](https://github.com/NovaSky-AI/SkyRL/pull/453)). +* Special thanks to the community contributors for this release: @HollowMan6, @fzyzcjy, @erictang000, @hawkoli1987. + +## NVIDIA Megatron-Bridge 0.1.0rc4 + +* Fix docs build +* Update performance scripts + +## NVIDIA Megatron-Bridge 0.1.0rc3 + +* Model Collection Support + * Llama + * Qwen 2, Qwen 3, Qwen 3 MoE + * DeepSeek + * Mamba +* [Migration guide from NeMo 2 to Megatron-Bridge](https://docs.nvidia.com/nemo/megatron-bridge/0.1.0/nemo2-migration-guide.html) +* [Contribution guide for adding a new model](https://docs.nvidia.com/nemo/megatron-bridge/0.1.0/adding-new-models.html) +* [Checkpoint conversion from Hugging Face to Megatron](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/conversion) +* [Performance](https://docs.nvidia.com/nemo/megatron-bridge/0.1.0/performance-summary.html) + * MoE LLM + * Change the model to dropless with balanced gating + * Fusion of operators in router function + * Global permutation fusion with A2A dispatcher + * EP A2A communication overlap with computation in both 1F1B pipelining and non-pipelined training + * Precision-aware optimizer update to support BF16 states + * Megatron FSDP + * Migration from mcore FSDP to megatron FSDP + * Fusion of weight gradient copy to reduce-scatter communication buffer to WGRAD GEMM + * Removed redundant optimizer operations + * Use Zero1 (opt and master param sharding) in the replica domain of hybrid FSDP to further lower memory usage + * IB-SHARP support for the IB AllReduce of hybrid FSDP in a patch with NCCL2.28 + * MXFP8 + * Improved act grad all-gather overlap performance via userbuffer + * Parameter all-gather overlap with computation while the communication buffer sharing with reduce-scatter + * Fusion of MXFP8 scaling factor swizzling kernels + * Use PDL (Programmatic Dependent Launch) for quantization kernels to lower CPU overhead + * Others + * Full iteration cuda graph for dense model without pipelining + * Fusion of activation and cast fusion (currently tensor-wise scaling only) + * Store SwiGLU input in FP8 to save activation memory + +## NVIDIA Megatron-Bridge 0.1.0a0 + +* Llama and Qwen +* Pretrain/SFT +* PeFT +* Recipe structure with examples for plain python & NeMo Run usage diff --git a/fern/v0.2.0/pages/releases/known-issues.mdx b/fern/v0.2.0/pages/releases/known-issues.mdx new file mode 100644 index 0000000000..dba67f0e8f --- /dev/null +++ b/fern/v0.2.0/pages/releases/known-issues.mdx @@ -0,0 +1,15 @@ +--- +title: "Known Issues" +description: "" +--- + +This page lists known issues and limitations in the current release. + +## 25.11 + +- Deepseek V3 on H100 has an issue when using DeepEP and fails with `RuntimeError: DeepEP error: timeout (dispatch CPU)`. +- MODEL_TFLOP/s/GPU is printed as 0 to stdout for all Hybrid models, such as Nemotron-H 56B. + +## 25.09 + +- **Pretraining DeepSeek in subchannel FP8 precision is not working.** Pretraining DeepSeek with current scaling FP8 is a workaround, but MTP loss does not converge. diff --git a/fern/v0.2.0/pages/releases/software-versions.mdx b/fern/v0.2.0/pages/releases/software-versions.mdx new file mode 100644 index 0000000000..c4cf014b61 --- /dev/null +++ b/fern/v0.2.0/pages/releases/software-versions.mdx @@ -0,0 +1,52 @@ +--- +title: "Software Component Versions" +description: "" +--- + +## NeMo Framework 25.11 + +| Software Component | Version | +|-------------------|---------| +| PyTorch | 2.9.0a0 | +| Megatron Core | dev:0.15.0 | +| Transformer Engine | 2.9 | +| Megatron-Bridge | 0.2.0 | +| Megatron-FSDP | 0.2.0 | +| Export-Deploy | 0.3.0 | +| Evaluator | 0.2.0 | +| NeMo | 2.6.0 | +| NeMo Run | 0.7.0 | +| TRT-ModelOpt | 0.37.0 | +| NVRX | 0.4.1 | +| CUDA | 13.0.1 | +| cuDNN | 9.13.1.26 | +| TRT-LLM | 1.1.0a0 | + + +NVIDIA NeMo™ Framework Training container is built on top of NVIDIA Optimized Frameworks PyTorch 25.06 container: https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/index.html + + +## NeMo Framework 25.09 + +| Software Component | Version | +|-------------------|---------| +| PyTorch | 2.8.0a0+5228986c39.nv25.6 | +| Megatron Core | 0.14.0 | +| Transformer Engine | 2.7.0 | +| Megatron-Bridge | 0.1.0rc3 | +| Megatron-FSDP | 0.1.0 | +| Automodel | 0.1.0 | +| Export-Deploy | 0.2.0 | +| Eval | 0.1.0 | +| NeMo | 2.5.0 | +| NeMo Run | 0.6.0 | +| NeMo Curator | 1.0.0 | +| TRT-ModelOpt | 0.35.1 | +| NVRX | 0.4.1 | +| CUDA | 12.9.1 | +| cuDNN | 9.11.0.98 | +| TRT-LLM (supported but not included) | 1.0.0 | + + +NVIDIA NeMo™ Framework Training container is built on top of NVIDIA Optimized Frameworks PyTorch 25.06 container: https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/index.html + diff --git a/fern/v0.2.0/pages/training/README.mdx b/fern/v0.2.0/pages/training/README.mdx new file mode 100644 index 0000000000..e319707274 --- /dev/null +++ b/fern/v0.2.0/pages/training/README.mdx @@ -0,0 +1,124 @@ +--- +title: "Training and Customization" +description: "" +--- + +This directory contains comprehensive documentation for training and customizing models with Megatron Bridge. Learn how to configure training, optimize performance, and customize training workflows. + +## Quick Navigation + +### I want to + +**🚀 Get started with training** +→ Start with [Configuration Container Overview](/config-container-overview) to understand the training setup + +**⚙️ Configure training parameters** +→ See [Training Loop Settings](/training-loop-settings) and [Optimizer & Scheduler](/optimizer-scheduler) + +**📊 Monitor and profile training** +→ Check [Logging](/logging) and [Profiling](/profiling) guides + +**💾 Manage checkpoints** +→ Read [Checkpointing](/checkpointing) for saving and resuming training + +**⚡ Optimize performance** +→ Explore [Performance Guide](/../performance-guide) and [Performance Summary](/../performance-summary) + +**🔧 Customize training** +→ See [PEFT](/peft), [Distillation](/distillation), [Entry Points](/entry-points), and [Callbacks](/callbacks) + +## Core Training Documentation + +### Configuration and Setup + +| Document | Purpose | When to Read | +|----------|---------|--------------| +| **[Configuration Container Overview](/config-container-overview)** | Central configuration object for all training settings | First time setting up training | +| **[Entry Points](/entry-points)** | Training entry points and execution flow | Understanding how training starts | +| **[Training Loop Settings](/training-loop-settings)** | Training loop parameters and configuration | Configuring batch sizes, iterations, validation | + +### Optimization and Performance + +| Document | Purpose | When to Read | +|----------|---------|--------------| +| **[Optimizer & Scheduler](/optimizer-scheduler)** | Optimizer and learning rate scheduler configuration | Setting up optimization | +| **[Mixed Precision](/mixed-precision)** | Mixed precision training for memory efficiency | Reducing memory usage | +| **[Communication Overlap](/communication-overlap)** | Overlapping communication with computation | Optimizing distributed training | +| **[Attention Optimizations](/attention-optimizations)** | Optimizing attention mechanisms | Improving training speed | +| **[Activation Recomputation](/activation-recomputation)** | Gradient checkpointing strategies | Reducing memory footprint | +| **[CPU Offloading](/cpu-offloading)** | Offloading to CPU for memory management | Working with limited GPU memory | + +### Monitoring and Debugging + +| Document | Purpose | When to Read | +|----------|---------|--------------| +| **[Logging](/logging)** | Logging configuration and TensorBoard/WandB integration | Monitoring training progress | +| **[Profiling](/profiling)** | Performance profiling and analysis | Identifying bottlenecks | +| **[Resiliency](/resiliency)** | Handling failures and recovery | Building robust training pipelines | + +### Advanced Features + +| Document | Purpose | When to Read | +|----------|---------|--------------| +| **[PEFT](/peft)** | Parameter-Efficient Fine-Tuning (LoRA, etc.) | Fine-tuning with limited resources | +| **[Packed Sequences](/packed-sequences)** | Sequence packing for efficiency | Optimizing data loading | +| **[Distillation](/distillation)** | Knowledge distillation techniques | Transferring knowledge between models | +| **[Checkpointing](/checkpointing)** | Checkpoint saving, loading, and resuming | Managing training state | +| **[Callbacks](/callbacks)** | Inject custom logic into training loop | Custom logging, metrics, third-party integrations | + +## Training Workflow + +A typical training workflow involves: + +1. **Configure Training** - Set up `ConfigContainer` with model, data, and training parameters +2. **Prepare Data** - Configure dataset loading and preprocessing +3. **Set Optimization** - Configure optimizer, scheduler, and mixed precision +4. **Enable Monitoring** - Set up logging and profiling +5. **Configure Checkpointing** - Set up checkpoint saving and resuming +6. **Launch Training** - Start training with configured entry points +7. **Monitor Progress** - Track metrics via logging and profiling +8. **Resume if Needed** - Use checkpointing to resume from saved state + +## Related Documentation + +- **[Main Documentation Index](/../index)** - Return to main documentation +- **[Performance Guide](/../performance-guide)** - Comprehensive performance optimization guide +- **[Performance Summary](/../performance-summary)** - Quick performance reference +- **[Recipe Usage](/../recipe-usage)** - Using training recipes +- **[Parallelisms](/../parallelisms)** - Understanding distributed training strategies +- **[Bridge Guide](/../bridge-guide)** - Working with Hugging Face models + +## Common Training Scenarios + +### 🆕 First-Time Training Setup + +1. [Configuration Container Overview](/config-container-overview) - Understand the configuration system +2. [Entry Points](/entry-points) - Learn how to start training +3. [Training Loop Settings](/training-loop-settings) - Configure basic training parameters +4. [Logging](/logging) - Set up monitoring + +### ⚡ Performance Optimization + +1. [Performance Guide](/../performance-guide) - Comprehensive optimization strategies +2. [Mixed Precision](/mixed-precision) - Enable mixed precision training +3. [Communication Overlap](/communication-overlap) - Optimize distributed training +4. [Activation Recomputation](/activation-recomputation) - Reduce memory usage +5. [Profiling](/profiling) - Identify bottlenecks + +### 💾 Production Training + +1. [Checkpointing](/checkpointing) - Reliable checkpoint management +2. [Resiliency](/resiliency) - Handle failures gracefully +3. [Logging](/logging) - Comprehensive monitoring +4. [Profiling](/profiling) - Performance analysis + +### 🔧 Customization + +1. [PEFT](/peft) - Parameter-efficient fine-tuning +2. [Distillation](/distillation) - Knowledge distillation +3. [Entry Points](/entry-points) - Custom training workflows +4. [Callbacks](/callbacks) - Inject custom logic (third-party integrations) + +--- + +**Ready to start training?** Begin with [Configuration Container Overview](/config-container-overview) or return to the [main documentation](/../README). diff --git a/fern/v0.2.0/pages/training/activation-recomputation.mdx b/fern/v0.2.0/pages/training/activation-recomputation.mdx new file mode 100644 index 0000000000..9f2e8477c4 --- /dev/null +++ b/fern/v0.2.0/pages/training/activation-recomputation.mdx @@ -0,0 +1,148 @@ +--- +title: "Activation Recomputation" +description: "" +--- + +The input activations of network layers are stored in device memory and are used to compute gradients during back-propagation. When training a LLM with a long sequence length or a large micro-batch size, these input activations can quickly saturate device memory. Checkpointing a few activations and recomputing the rest is a common technique to reduce device memory usage. + +Activation recomputation in Megatron Bridge is configured through the model provider's recomputation parameters, which are based on Megatron Core's `TransformerConfig`. + +## Transformer Layer Recomputation + +Megatron Bridge supports transformer layer recomputation, which checkpoints the input of each transformer layer and recomputes the activations for the remaining layers. This technique significantly reduces activation memory usage. However, it increases the per-transformer layer computation cost by 30% due to re-executing the entire layer's forward computation. + +Megatron Bridge also supports partial transformer layer recomputation, which is beneficial when recomputing a few transformer layers helps to reduce enough GPU memory for the model to fit. This approach avoids the need to recompute the rest of the layers. + +### Configuration + +Transformer layer recomputation is configured through the model provider's recomputation parameters: + +```python +from megatron.bridge.models import GPTModelProvider + +# Full recomputation - recompute all layers +model_config = GPTModelProvider( + recompute_granularity="full", # Enable full layer recomputation + recompute_method="uniform", # Uniform distribution across layers + recompute_num_layers=4, # Number of layers per recomputation block + # ... other model parameters +) +``` + +### Recomputation Methods + +#### Block Method +Recomputes a specific number of transformer layers per pipeline stage: + +```python +model_config = GPTModelProvider( + recompute_granularity="full", + recompute_method="block", # Block-wise recomputation + recompute_num_layers=4, # Recompute 4 layers per pipeline stage +) +``` + +#### Uniform Method +Uniformly divides the total number of transformer layers and recomputes input activations for each divided chunk: + +```python +model_config = GPTModelProvider( + recompute_granularity="full", + recompute_method="uniform", # Uniform distribution + recompute_num_layers=8, # Number of layers per recomputation block +) +``` + +### Pipeline Parallelism Considerations + +When training with pipeline parallelism: +- `recompute_num_layers` indicates the layers per pipeline stage +- When using virtual pipelining, `recompute_num_layers` specifies the number of layers per virtual pipeline stage +- The framework automatically handles recomputation coordination across pipeline stages + +![Activation Recomputation Methods](/assets/training/images/activation-recomputation-example-1.jpg) +*Figure 1: Scheme of uniform and block checkpointing method (full checkpointing granularity)* + +## Self-attention Recomputation + +Megatron Bridge supports selective self-attention recomputation that checkpoints the inputs of each self-attention block and recomputes the intermediate input activations. This cost-efficient method achieves high memory savings with minimal recomputation cost. + +The intermediate layers of the self-attention block account for the majority of the activation memory because the input sizes of softmax, dropout, and QKV dot-product attention layers have memory complexity proportional to the sequence length squared. However, their recomputation cost is relatively smaller than other linear projection layers that scale with the hidden size squared. + +![Activation Recomputation Granularity](/assets/training/images/activation-recomputation-example-2.jpg) +*Figure 2: Scheme of full and selective checkpointing granularity* + +### Configuration + +Self-attention recomputation is enabled using selective granularity: + +```python +from megatron.bridge.models import GPTModelProvider + +model_config = GPTModelProvider( + recompute_granularity="selective", # Enable selective recomputation + recompute_modules=["core_attn"], # Recompute attention modules (default) + # ... other model parameters +) +``` + +### Recomputation Modules + +Megatron Bridge supports selective recomputation for various modules: + +```python +model_config = GPTModelProvider( + recompute_granularity="selective", + recompute_modules=[ + "core_attn", # Core attention computation (default) + "mlp", # MLP layers + "layernorm", # Layer normalization + "moe", # Mixture of Experts layers + "moe_act", # MoE activation functions + "shared_experts", # Shared expert layers + "mla_up_proj", # Multi-Latent Attention up projection + ], +) +``` + +### Flash Attention Integration + +Self-attention recomputation is automatically enabled when using Flash Attention through Transformer Engine. Flash Attention inherently provides memory efficiency by recomputing attention scores rather than storing them, making additional explicit recomputation often unnecessary. + +## Advanced Recomputation Configuration + +### Distributed Activation Checkpointing + +For models using model parallelism, you can distribute saved activations across the model parallel group: + +```python +model_config = GPTModelProvider( + recompute_granularity="selective", + distribute_saved_activations=True, # Distribute across model parallel group + # Note: Cannot be used with sequence_parallel=True +) +``` + +### Memory vs Computation Trade-offs + +Different recomputation strategies offer different memory-computation trade-offs: + +- **Selective recomputation**: Provides high memory savings with minimal recomputation cost by targeting memory-intensive operations like attention +- **Full recomputation**: Significantly reduces activation memory usage but increases per-transformer layer computation cost by approximately 30% +- **No recomputation**: Preserves all activations in memory, requiring more GPU memory but no additional computation + +### MoE-Specific Recomputation + +For Mixture of Experts models, specialized recomputation options are available: + +```python +model_config = GPTModelProvider( + # MoE configuration + num_moe_experts=8, + expert_model_parallel_size=2, + + # MoE recomputation + recompute_granularity="selective", + recompute_modules=["moe", "moe_act"], # Recompute MoE-specific modules +) +``` diff --git a/fern/v0.2.0/pages/training/attention-optimizations.mdx b/fern/v0.2.0/pages/training/attention-optimizations.mdx new file mode 100644 index 0000000000..3114690eb8 --- /dev/null +++ b/fern/v0.2.0/pages/training/attention-optimizations.mdx @@ -0,0 +1,131 @@ +--- +title: "Attention Optimizations" +description: "" +--- + +Megatron Bridge provides several attention optimizations to improve the efficiency and performance of transformer models. These optimizations include Flash Attention for memory efficiency, and Multi-Query Attention (MQA) and Grouped-Query Attention (GQA) for computational efficiency. + +## Flash Attention + +### Overview + +Flash attention is an algorithm designed to improve the efficiency of the attention mechanism in transformer models such as GPT and BERT. The attention mechanism has quadratic time and memory complexity in sequence length and can present significant runtime and memory challenges for longer sequences. + +Compared to the standard, non-flash algorithm, flash attention applies two techniques to lower the memory requirement and improve compute efficiency: + +1. **Tiling technique**: Decomposes the inputs based on the shared memory size and calculates the softmax one tile at a time. Instead of working on the entire query, key, and value tensors at once, it makes several passes at these tensors and then combines the results in a subsequent step. + +2. **Recomputation technique**: Stores the softmax normalization factors (linear to sequence length), instead of the softmax results (quadratic to sequence length), and uses these normalization factors to recompute the attention scores. This saves the amount of data to write to global memory and reduces both the I/O traffic between global memory and shared memory. + +Flash attention lowers the memory footprint and computational complexity from quadratic to linear, greatly extending the range of sequence length allowed in large language models. + +### Configure Flash Attention + +In Megatron Bridge, flash attention is configured through the `attention_backend` parameter in your model configuration. The framework supports multiple attention backends through Transformer Engine integration: + +```python +from megatron.bridge.models import GPTModelProvider +from megatron.core.transformer.enums import AttnBackend + +# Configure model with flash attention (default) +model_config = GPTModelProvider( + attention_backend=AttnBackend.auto, # Let TE choose the best backend (default) + # ... other model parameters +) + +# Or explicitly specify flash attention +model_config = GPTModelProvider( + attention_backend=AttnBackend.flash_attn, # Explicitly use flash attention + # ... other model parameters +) +``` + +### Attention Backend Options + +Megatron Bridge supports several attention backends through the `attention_backend` configuration: + +- `AttnBackend.auto`: Automatically selects the best available backend (recommended) +- `AttnBackend.flash_attn`: Explicitly use Flash Attention implementation +- `AttnBackend.fused_attn`: Use cuDNN fused attention (when available) +- `AttnBackend.local`: Use local PyTorch implementation (for debugging) + +### Environment Variable Control + +For fine-grained control, you can still use environment variables to disable specific implementations: + +```bash +# Disable flash attention +export NVTE_FLASH_ATTN=0 + +# Disable cuDNN flash attention +export NVTE_FUSED_ATTN=0 +``` + +However, the recommended approach is to use the `attention_backend` configuration parameter. + +## Multi-query Attention (MQA) and Grouped-query Attention (GQA) + +**Multi-query Attention (MQA)** and **Grouped-query Attention (GQA)** are modifications of the traditional multihead attention mechanism in Transformer models. These methods improve the efficiency and effectiveness of attention mechanisms. + +### Overview + +**Multi-query Attention (MQA)** + +MQA treats all attention heads as a single group, reducing computational complexity and accelerating training times. It is beneficial when model scalability or limited computational resources are concerns. + +**Grouped-query Attention (GQA)** + +GQA groups the heads into clusters, each processing a subset of queries independently. This method balances the detailed focus of traditional multihead attention with the broad approach of MQA, enhancing nuanced input data processing. + +These attention variants offer: + +- **Reduced computational load**: Both methods decrease computation, beneficial for large models +- **Increased processing speed**: Simplifying attention leads to faster training and inference +- **Flexibility and adaptability**: Adjustments can be made based on task needs or hardware constraints + +### Enable MQA and GQA + +To use MQA or GQA in Megatron Bridge, adjust the `num_query_groups` parameter in your model configuration: + +#### Multi-query Attention (MQA) +Set `num_query_groups` to 1 to treat all attention heads as a single group: + +```python +from megatron.bridge.models import GPTModelProvider + +model_config = GPTModelProvider( + num_attention_heads=32, + num_query_groups=1, # Enables Multi-query Attention + # ... other model parameters +) +``` + +#### Grouped-query Attention (GQA) +Set `num_query_groups` to a number that is a divisor of the total number of attention heads (more than one but less than the total heads): + +```python +model_config = GPTModelProvider( + num_attention_heads=32, + num_query_groups=8, # Enables Grouped-query Attention (4 heads per group) + # ... other model parameters +) +``` + +#### Regular Multihead Attention +For regular attention, set this parameter to `None` or match it with the number of heads: + +```python +model_config = GPTModelProvider( + num_attention_heads=32, + num_query_groups=None, # Default setting for regular multihead attention + # Or equivalently: + # num_query_groups=32, # One group per head + # ... other model parameters +) +``` + +## Resources + +- [Megatron Core Attention Implementation](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/transformer/attention.py) +- [Flash Attention Paper](https://arxiv.org/abs/2205.14135) +- [Transformer Engine Attention Mechanisms](https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/examples/attention/attention.html) diff --git a/fern/v0.2.0/pages/training/callbacks.mdx b/fern/v0.2.0/pages/training/callbacks.mdx new file mode 100644 index 0000000000..14be1a0ee3 --- /dev/null +++ b/fern/v0.2.0/pages/training/callbacks.mdx @@ -0,0 +1,246 @@ +--- +title: "Callbacks" +description: "" +--- + +Megatron Bridge provides a lightweight callback system for injecting custom logic into the training and evaluation loop without modifying framework code. This is ideal for propietary integrations or custom logging and metrics tracking. + +## Quick Start + +### Class-Based Callbacks + +Subclass [`bridge.training.callbacks.Callback`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.callbacks.Callback.html) and override event methods: + +```python +import time + +from megatron.bridge.training.callbacks import Callback +from megatron.bridge.training.gpt_step import forward_step +from megatron.bridge.training.pretrain import pretrain +from megatron.bridge.recipes.qwen import qwen25_500m_pretrain_config + +class MyCallback(Callback): + def on_train_start(self, context): + context.user_state['start_time'] = time.time() + print(f"Training started at step {context.state.train_state.step}") + + def on_train_step_end(self, context): + if context.loss_dict: + print(f"Step {context.state.train_state.step}: loss={context.loss_dict}") + + def on_train_end(self, context): + elapsed = time.time() - context.user_state['start_time'] + print(f"Training completed in {elapsed:.2f}s") + +# Create a config that fits on a single GPU +config = qwen25_500m_pretrain_config() + +# Pass callbacks to pretrain +pretrain(config, forward_step, callbacks=[MyCallback()]) +``` + +### Functional Callbacks + +Register functions directly with [`bridge.training.callbacks.CallbackManager`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.callbacks.CallbackManager.html): + +```python +from megatron.bridge.training.callbacks import CallbackManager +from megatron.bridge.training.gpt_step import forward_step +from megatron.bridge.training.pretrain import pretrain +from megatron.bridge.recipes.qwen import qwen25_500m_pretrain_config + +def log_step(context): + step = context.state.train_state.step + if context.loss_dict: + print(f"Step {step}: {context.loss_dict}") + +callback_manager = CallbackManager() +callback_manager.register("on_train_step_end", log_step) + +# Create a config that fits on a single GPU +config = qwen25_500m_pretrain_config() + +pretrain(config, forward_step, callbacks=callback_manager) +``` + +### Mixing Both Patterns + +Both registration patterns can be combined: + +```python +from megatron.bridge.training.callbacks import CallbackManager +from megatron.bridge.training.gpt_step import forward_step +from megatron.bridge.training.pretrain import pretrain +from megatron.bridge.recipes.qwen import qwen25_500m_pretrain_config + +manager = CallbackManager() +manager.add(MyCallback()) +manager.add([TimingCallback(), MetricsCallback()]) +manager.register("on_eval_end", lambda ctx: print("Evaluation complete!")) + +# Create a config that fits on a single GPU +config = qwen25_500m_pretrain_config() + +pretrain(config, forward_step, callbacks=manager) +``` + +## Available Events + +### Training Events + +| Event | When Fired | Available Context Fields | +|-------|------------|-------------------------| +| `on_train_start` | After `model.train()`, before training loop | `state`, `model`, `user_state`, `optimizer`, `scheduler` | +| `on_train_step_start` | Before each training step | `state`, `model`, `user_state`, `optimizer`, `scheduler` | +| `on_train_step_end` | After each training step | `state`, `model`, `user_state`, `optimizer`, `scheduler`, `loss_dict`, `grad_norm`, `skipped_iter` | +| `on_train_end` | After training loop completes | `state`, `model`, `user_state`, `optimizer`, `scheduler` | + +### Validation Events + +| Event | When Fired | Available Context Fields | +|-------|------------|-------------------------| +| `on_eval_start` | After `model.eval()`, before validation loop | `state`, `model`, `user_state` | +| `on_eval_step_start` | Before each validation step | `state`, `model`, `user_state` | +| `on_eval_step_end` | After each validation step | `state`, `model`, `user_state` | +| `on_eval_end` | After validation completes | `state`, `model`, `user_state`, `total_loss_dict` | + +### Test Events + +| Event | When Fired | Available Context Fields | +|-------|------------|-------------------------| +| `on_test_start` | After `model.eval()`, before test loop | `state`, `model`, `user_state` | +| `on_test_step_start` | Before each test step | `state`, `model`, `user_state` | +| `on_test_step_end` | After each test step | `state`, `model`, `user_state` | +| `on_test_end` | After test completes | `state`, `model`, `user_state`, `total_loss_dict` | + +## CallbackContext + +The [`bridge.training.callbacks.CallbackContext`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.callbacks.CallbackContext.html) provides access to framework state: + +### Always Available + +- **`state`**: [`bridge.training.state.GlobalState`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.state.GlobalState.html) - Contains config, train_state, timers, and loggers +- **`model`**: List of model chunks +- **`user_state`**: Mutable dict for storing data across callback invocations + +### Training Events Only + +- **`optimizer`**: The optimizer instance +- **`scheduler`**: Learning rate scheduler + +### Event-Specific Fields + +- **`loss_dict`** (`on_train_step_end`): Dictionary of reduced losses from the training step +- **`grad_norm`** (`on_train_step_end`): Gradient norm (if computed) +- **`skipped_iter`** (`on_train_step_end`): Whether the iteration was skipped +- **`total_loss_dict`** (`on_eval_end`, `on_test_end`): Aggregated evaluation/test losses + +## User State + +The `CallbackManager` owns a `user_state` dictionary that persists across all callback invocations during a training run. Use it to share data between callbacks or accumulate metrics: + +```python +class StepCounterCallback(Callback): + def on_train_start(self, context): + context.user_state['callback_step_count'] = 0 + + def on_train_step_end(self, context): + context.user_state['callback_step_count'] += 1 + + def on_train_end(self, context): + print(f"Callback saw {context.user_state['callback_step_count']} steps") +``` + +## Distributed Training + +Callbacks fire on **all ranks** without framework-level synchronization. If your callback should only run on specific ranks, add guards: + +```python +import torch.distributed as dist + +class RankZeroCallback(Callback): + def on_train_step_end(self, context): + if dist.get_rank() == 0: + print(f"Step {context.state.train_state.step} complete") +``` + +## Exception Handling + +Exceptions from callbacks propagate to the caller. The framework does not catch or handle callback exceptions. If your callback might fail, wrap it in a try-except: + +```python +def safe_callback(context): + try: + # Your logic here + external_service.log(context.loss_dict) + except Exception as e: + print(f"Callback failed: {e}") + # Don't re-raise to avoid stopping training +``` + +## Execution Order + +Callbacks fire in registration order: + +1. Callbacks added via `add()` fire in the order they were added +2. Callbacks registered via `register()` fire in the order they were registered +3. If both methods are used, the order depends on when each was called + +## Introspection + +Query registered callbacks: + +```python +manager = CallbackManager() +manager.register("on_train_start", my_fn) + +# Check if any callbacks exist for an event +if manager.has_callbacks("on_train_start"): + print("Callbacks registered for on_train_start") + +# List all callbacks for an event +callbacks = manager.list_callbacks("on_train_start") +print(f"Found {len(callbacks)} callbacks") + +# Get all valid event names +print(manager.events) # frozenset of valid event names +``` + +## Design Principles + +The callback system follows these principles: + +1. **First-Party Isolation**: Framework code never uses callbacks for its own logic. Callbacks are strictly for third-party extensions. + +2. **Zero Overhead**: When no callbacks are registered, there is zero performance overhead. + +3. **Safety**: Callbacks receive framework state but modifying it is at the user's own risk. The framework makes no guarantees about the effects of modifications. + +## Examples + +### Proprietary Metrics + +```python +class ProprietaryMetricsCallback(Callback): + """Send metrics to internal monitoring system.""" + + def __init__(self, endpoint: str): + self.client = InternalMetricsClient(endpoint) + + def on_train_step_end(self, context): + if context.loss_dict: + self.client.send({ + "step": context.state.train_state.step, + "loss": context.loss_dict.get("lm loss"), + "grad_norm": context.grad_norm, + "cluster_id": os.environ.get("CLUSTER_ID"), + }) +``` + +## API Reference + +- [`bridge.training.callbacks.Callback`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.callbacks.Callback.html) +- [`bridge.training.callbacks.CallbackContext`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.callbacks.CallbackContext.html) +- [`bridge.training.callbacks.CallbackManager`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.callbacks.CallbackManager.html) +- [`bridge.training.callbacks.normalize_callbacks`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.callbacks.normalize_callbacks.html) +- [`bridge.training.callbacks.should_fire`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.callbacks.should_fire.html) diff --git a/fern/v0.2.0/pages/training/checkpointing.mdx b/fern/v0.2.0/pages/training/checkpointing.mdx new file mode 100644 index 0000000000..a55defe86f --- /dev/null +++ b/fern/v0.2.0/pages/training/checkpointing.mdx @@ -0,0 +1,227 @@ +--- +title: "Checkpointing" +description: "" +--- + +The [`bridge.training.config.CheckpointConfig`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.CheckpointConfig.html) controls model checkpointing behavior, including saving and loading checkpoints, checkpoint formats, and various optimization features. + + +This documentation covers **Megatron-format checkpoints** used during training. For converting between 🤗 Hugging Face and Megatron formats, see the [Bridge Guide](/bridge-guide). + + +## Overview + +Megatron Bridge uses Megatron Core's distributed checkpointing system, which is designed for large-scale training across multiple GPUs and nodes. The distributed checkpoint approach saves the state of a distributed training job by sharding checkpoint data across multiple files, reducing memory overhead and improving GPU utilization during save/load operations. + +### Distributed Checkpointing Benefits + +**Memory Efficiency**: Instead of gathering all model parameters and optimizer states on a single rank, distributed checkpointing saves data directly from each rank, significantly reducing memory requirements during checkpointing. + +**Parallelism Flexibility**: The system provides flexibility to resume training using different parallelism strategies. You can change tensor parallelism, pipeline parallelism, or data parallelism sizes between checkpoint save and load operations. + +**Scalability**: Handles all types of parallelism including: +- **Data Parallelism (DP)**: Replicates the model across multiple GPUs with different data batches +- **Tensor Parallelism (TP)**: Distributes individual layer parameters across GPUs +- **Pipeline Parallelism (PP)**: Assigns consecutive layers to different GPUs +- **Context Parallelism (CP)**: Shards tensors along the sequence dimension for long sequences +- **Expert Parallelism (EP)**: Distributes MoE expert weights across GPUs + +**Performance**: The distributed optimizer shards optimizer states and master parameters across data-parallel ranks instead of replicating them, reducing memory usage and communication overhead. + +## Save Configuration + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `save` | `Optional[str]` | `None` | Output directory to save checkpoints to **in Megatron format** | +| `save_interval` | `Optional[int]` | `None` | Number of iterations between persistent checkpoint saves | +| `save_optim` | `bool` | `True` | Whether to save optimizer state | +| `save_rng` | `bool` | `True` | Whether to save random number generator state | +| `save_tokenizer_assets` | `bool` | `True` | Whether to save tokenizer files (vocab, config, special tokens) to checkpoint | + +### Asynchronous Saving + +Asynchronous saving allows training to continue while checkpoint data is persisted to disk in the background, reducing the impact of checkpointing on training throughput. + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `async_save` | `bool` | `False` | Enable asynchronous checkpoint saving (requires `torch_dist` format) | + +## Load Configuration + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `load` | `Optional[str]` | `None` | Directory containing a model checkpoint to load **in Megatron format** | +| `load_optim` | `bool` | `True` | Whether to load optimizer state from checkpoint | +| `load_rng` | `bool` | `True` | Whether to load random number generator state from checkpoint | +| `load_main_params_from_ckpt` | `bool` | `False` | Load main parameters from checkpoint (use with `load_optim=False`) | +| `ckpt_step` | `Optional[int]` | `None` | Specific checkpoint iteration to load (overrides latest from tracker) | +| `exit_on_missing_checkpoint` | `bool` | `False` | Exit if specified checkpoint is not found instead of random initialization | +| `dist_ckpt_strictness` | `Literal[...]` | `"assume_ok_unexpected"` | Handling of key mismatches during distributed checkpoint load | + +### Loading Specific Checkpoint Iterations + +By default, Megatron Bridge loads the **latest checkpoint** available in the specified directory by reading from the tracker file (`latest_train_state.pt`). However, you can explicitly load from a specific checkpoint iteration using the `ckpt_step` parameter. + +**Python API:** +```python +from megatron.bridge.training.config import CheckpointConfig + +# Load latest checkpoint +checkpoint = CheckpointConfig( + load="/path/to/checkpoint_dir" +) + +# Load specific iteration +checkpoint = CheckpointConfig( + load="/path/to/checkpoint_dir", + ckpt_step=5000 # Overrides tracker, loads iter_0005000 +) +``` + + +The `load` parameter should always point to the base checkpoint directory (not the `iter_N` subdirectory). The `ckpt_step` parameter overrides which iteration is loaded from that directory. + +**Important:** If `ckpt_step` is specified but the checkpoint directory does not exist, training will **fail immediately** with a `FileNotFoundError`. This is intentional to prevent accidentally starting training from scratch when you meant to resume from a specific checkpoint. + +**PEFT Note:** The `ckpt_step` parameter applies **only to the `load` path** (adapter checkpoints), not to `pretrained_checkpoint` (frozen base model). When resuming PEFT training: +- `pretrained_checkpoint`: Always loads the latest/release checkpoint (base model) +- `load` + `ckpt_step`: Can load a specific adapter checkpoint iteration + +### Checkpoint Loading Strictness + +When loading distributed checkpoints, there may be mismatches between the keys in the saved checkpoint and what the current model expects. This can happen when resuming training with different parallelism settings, model configurations, or software versions. The `dist_ckpt_strictness` parameter controls how these mismatches are handled: + +- **`assume_ok_unexpected`**: Assume unexpected keys are acceptable (default, most permissive) +- **`log_unexpected`**: Log unexpected keys but continue loading +- **`log_all`**: Log all key mismatches for debugging +- **`raise_unexpected`**: Raise error on unexpected keys (stricter validation) +- **`raise_all`**: Raise error on any key mismatch (strictest validation) +- **`return_unexpected`**: Return information about unexpected keys +- **`return_all`**: Return information about all key mismatches +- **`ignore_all`**: Ignore all key mismatches completely + +## Fine-tuning Configuration + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `pretrained_checkpoint` | `Optional[str]` | `None` | Directory containing pretrained model checkpoint **in Megatron format** for fine-tuning | + +## Checkpoint Format + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `ckpt_format` | `Literal["torch_dist"]` | `"torch_dist"` | Checkpoint format (PyTorch distributed checkpoint format) | + +## Performance Optimizations + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `fully_parallel_save` | `bool` | `True` | Apply full save parallelization across data parallel ranks | +| `fully_parallel_load` | `bool` | `False` | Apply full load parallelization across data parallel ranks | +| `ckpt_assume_constant_structure` | `bool` | `False` | Assume constant model/optimizer structure over successive checkpoint saves for performance optimizations | + +## Checkpoint Contents + +The checkpoint includes the following components when using the `torch_dist` checkpoint format: +- **Model parameters and optimizer states**: Stored across `.distcp` files to support distributed training. +- **Training state**: Captures the current iteration count, number of consumed samples, and the state of the learning rate scheduler. +- **Configuration**: Serialized as a YAML file (`run_config.yaml`) containing the complete `ConfigContainer`. +- **Tokenizer files**: All tokenizer artifacts (vocabulary, special tokens, config) for self-contained checkpoints. +- **Dataloader states**: Ensures deterministic resumption of data iteration. +- **Metadata**: Used for validating and correctly loading the checkpoint. + +Megatron Bridge creates checkpoints with the following directory structure: + + +checkpoint_dir/ +├── latest_train_state.pt # Latest training state (top-level) +├── iter_N/ # Checkpoint at iteration N +│ ├── __0_0.distcp # Distributed checkpoint shards: maps to PyTorch DCP weights format +│ ├── __0_1.distcp # Contains model parameters, optimizer states +│ ├── __1_0.distcp +│ ├── __1_1.distcp +│ ├── ... +│ ├── .metadata # PyTorch DCP checkpoint metadata +│ ├── common.pt # MCore dist ckpt states saved from rank 0 +│ ├── metadata.json # MCore dist ckpt metadata +│ ├── run_config.yaml # Serialized ConfigContainer +│ ├── train_state.pt # Number of steps, consumed samples, etc +│ ├── tokenizer/ # Tokenizer files (saved by default) +│ │ ├── tokenizer.json # Full tokenizer vocabulary +│ │ ├── tokenizer_config.json # Tokenizer configuration +│ │ ├── special_tokens_map.json # Special token definitions +│ │ └── ... # Other tokenizer artifacts +│ ├── dataloader_state/ # Data iterator states +│ │ ├── train_dataloader_dprank000.pt # DP rank 0 dataloader state +│ │ ├── train_dataloader_dprank001.pt # DP rank 1 dataloader state +│ │ ├── train_dataloader_dprank002.pt # DP rank 2 dataloader state +│ │ └── ... # One file per DP rank +``` + +### Tokenizer Assets + +By default, Megatron Bridge saves all tokenizer files to the checkpoint directory, making checkpoints self-contained and portable. This is particularly important for: +- **Inference and evaluation**: Direct access to tokenizer for computing logprobs +- **Portability**: No dependency on original tokenizer file locations +- **Reproducibility**: Exact tokenizer state is preserved + +The tokenizer files saved depend on the tokenizer type: +- **HuggingFace tokenizers**: `tokenizer.json`, `tokenizer_config.json`, `special_tokens_map.json`, and vocab files +- **SentencePiece tokenizers**: `tokenizer.model` file +- **GPT2 BPE tokenizers**: `vocab.json` and `merges.txt` +- **BERT tokenizers**: `vocab.txt` +- **Tiktoken tokenizers**: `tokenizer.json` + +To disable tokenizer asset saving for performance-sensitive scenarios: + +```python +from megatron.bridge.training.config import CheckpointConfig + +checkpoint = CheckpointConfig( + save_tokenizer_assets=False, # Skip tokenizer file saving + ... +) +``` + +Or in YAML: + +```yaml +checkpoint: + save_tokenizer_assets: false +``` + +## Local Checkpointing + +Local checkpointing saves model checkpoints directly to storage on each node (e.g., local SSDs or RAM disks), instead of relying solely on a shared network filesystem. This approach can significantly speed up the saving process and reduce the load on shared storage infrastructure. + +Local checkpointing leverages the [NVIDIA Resiliency Extension](https://nvidia.github.io/nvidia-resiliency-ext/checkpointing/local/index.html) and provides several key features: + +- **Local Saving**: Each node saves its part of the checkpoint locally, reducing network I/O and improving save performance. +- **Synchronous and Asynchronous Support**: Saving can happen synchronously or asynchronously, mirroring the configuration used for global checkpoints. +- **Automatic Cleanup**: Handles the removal of outdated or incomplete local checkpoints automatically. +- **Optional Replication**: For multi-node jobs, checkpoints are replicated to other nodes to allow recovery even if a node fails after saving. Single-node jobs do not use replication. +- **Automated Loading**: When resuming, the framework automatically finds the latest valid checkpoint, comparing local and global checkpoints, and retrieves any needed parts across nodes. +### Non-Persistent Checkpointing Configuration + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `non_persistent_save_interval` | `Optional[int]` | `None` | Iterations between non-persistent saves | +| `non_persistent_ckpt_type` | `Optional[Literal["global", "local", "in_memory", "None"]]` | `None` | Type of non-persistent checkpointing | +| `non_persistent_global_ckpt_dir` | `Optional[str]` | `None` | Directory for global non-persistent checkpoints | +| `non_persistent_local_ckpt_dir` | `Optional[str]` | `None` | Directory for local non-persistent checkpoints | +| `non_persistent_local_ckpt_algo` | `Literal["fully_parallel", "atomic"]` | `"fully_parallel"` | Algorithm for local non-persistent checkpointing | + +### Replication and Fault Tolerance + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `replication` | `bool` | `False` | Enable replication of local checkpoints across ranks | +| `replication_jump` | `Optional[int]` | `None` | Spacing between ranks storing replicas | +| `replication_factor` | `int` | `2` | Number of machines storing replica of each rank's data | + +### Checkpointing Distributed Optimizer + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `dist_ckpt_optim_fully_reshardable` | `bool` | `False` | Make optimizer distributed checkpoint fully reshardable (TP/PP/EP/DP) as opposed to plain DP reshardability | +| `distrib_optim_fully_reshardable_mem_efficient` | `bool` | `False` | Use as little memory as possible during save and load by using Gloo. Has affect only with `dist_ckpt_optim_fully_reshardable` flag | diff --git a/fern/v0.2.0/pages/training/communication-overlap.mdx b/fern/v0.2.0/pages/training/communication-overlap.mdx new file mode 100644 index 0000000000..55982345e6 --- /dev/null +++ b/fern/v0.2.0/pages/training/communication-overlap.mdx @@ -0,0 +1,234 @@ +--- +title: "Communication Overlap" +description: "" +--- + +Megatron Bridge supports overlapping communication with computation in distributed training to improve performance and throughput. This optimization technique reduces the impact of inter-GPU communication overhead by executing communication operations concurrently with computational operations whenever possible. + +Communication overlap is managed through the [`bridge.training.comm_overlap.CommOverlapConfig`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.comm_overlap.CommOverlapConfig.html) class and can be applied to different types of parallelism: tensor parallelism (TP), pipeline parallelism (PP), data parallelism (DP), and context parallelism (CP). + +## Data-parallel Communication Overlap + +Megatron Bridge supports the overlap of data-parallel (DP) communications with computations in LLM training. The framework features a Distributed Optimizer that distributes optimizer states and high-precision master parameters across GPUs. This introduces two types of data-parallel communications: reduce-scatter of gradients and all-gather of updated parameters. + +The DP communication is chunked by the granularity of a Transformer layer and overlaps each communication chunk with computation. This overlap method exposes only one DP communication chunk ensuring efficient large-scale LLM training. When training with pipeline parallelism, the granularity of DP communication becomes the Transformer layers per virtual pipeline stage. + +### Configuration + +DP communication overlap settings can be inspected in Megatron Core via the `DistributedDataParallelConfig` class. DP gradient reduce-scatter and parameter all-gather overlaps are enabled when setting `overlap_grad_reduce=True` and `overlap_param_gather=True`, respectively. The precision of gradient reduce-scatter is controlled by `grad_reduce_in_fp32`. When `grad_reduce_in_fp32=False`, gradients are reduced in bf16, leading to improved performance in large-scale training compared to the default fp32 precision. When training in fp8 computing precision, setting `fp8_param_gather=True` conducts the parameter all-gather in fp8, reducing the all-gather overhead by half. + +Data parallel communication overlap settings are controlled through the distributed data parallel and communication overlap configurations. + + +Data-parallel overlap relies on attributes such as `grad_reduce_in_fp32` and `fp8_param_gather`. When a mixed-precision recipe (for example `bf16_mixed`, `fp16_mixed`, `bf16_with_fp8_delayed_scaling_mixed`, etc.) is provided, those attributes are sourced from the recipe stored in the `MixedPrecisionConfig`. Set the desired values inside the mixed-precision configuration rather than overriding them directly on the optimizer or DDP configs. This ensures the communication overlap settings and the selected precision recipe remain consistent. + + +For example: + +```python +from megatron.bridge.training.config import ConfigContainer, OptimizerConfig +from megatron.bridge.training.comm_overlap import CommOverlapConfig +from megatron.bridge.training.mixed_precision import get_mixed_precision_config + +# Configure communication overlap +comm_overlap_config = CommOverlapConfig( + tp_comm_overlap=False, # Tensor parallel overlap + overlap_grad_reduce=True, # Gradient reduce-scatter overlap + overlap_param_gather=True, # Parameter all-gather overlap + overlap_param_gather_with_optimizer_step=False, # Advanced optimization + bucket_size=128 * 1024 * 1024, # 128MB bucket size +) + +# Configure distributed optimizer +optimizer_config = OptimizerConfig( + optimizer="adam", + lr=3e-4, + use_distributed_optimizer=True, # Required for DP overlap + # ... other optimizer parameters +) + +# Mixed precision configuration controls overlap-related attributes +mixed_precision_config = get_mixed_precision_config("bf16_mixed") +mixed_precision_config.grad_reduce_in_fp32 = False # Use bf16 for gradient reduction +mixed_precision_config.fp8_param_gather = False + +config = ConfigContainer( + comm_overlap=comm_overlap_config, + optimizer=optimizer_config, + mixed_precision=mixed_precision_config, + # ... other config parameters +) +``` + +Key data parallel overlap options: + +- `overlap_grad_reduce`: Overlaps gradient reduce-scatter with computation (default: True) +- `overlap_param_gather`: Overlaps parameter all-gather with computation (default: True) +- `overlap_param_gather_with_optimizer_step`: Advanced optimization for pipeline parallelism +- `bucket_size`: Controls the granularity of communication chunking (default: 128MB) +- `grad_reduce_in_fp32`: Controls gradient reduction precision (False for bf16, True for fp32) +- `fp8_param_gather`: Enables fp8 parameter all-gather for reduced communication overhead + +## Tensor-parallel Communication Overlap + +Tensor parallelism, used with sequence-parallel activation sharding (`sequence_parallel=True`), introduces activation (gradient) all-gather and reduce-scatter operations. Megatron Bridge provides various options to overlap the tensor-parallel (TP) communications with computation. + +![Tensor-parallel Communication Overlap](/assets/training/images/tp_comm_overlap.png) +*Figure: Tensor-parallel communication overlap showing bulk and pipelined overlap strategies.* + +The TP communication without direct computation dependency are overlapped with the computation in bulk (the linear layer and TP communication pairs in the yellow boxes). The bulk TP communication is enabled by default. The other TP communications with direct computation dependency are overlapped in pipelined fashion (the linear layer and TP communication pairs in the red boxes). + +In the pipelined overlap, the activation (gradient) tensor all-gather is replaced with multiple steps of input P2P ring exchanges, and reduce-scatter is replaced with multiple steps of GEMM output P2P ring exchanges followed by a reduction of the received outputs. + +### Configuration + +```python +from megatron.bridge.training.comm_overlap import ( + CommOverlapConfig, + TransformerLayerTPOverlapCfg, + userbuffers_bf16_h100_h8192_tp4_mbs1_seqlen8192 +) + +# Configure tensor parallel overlap +comm_overlap_config = CommOverlapConfig( + tp_comm_overlap=True, # Enable TP communication overlap + tp_comm_overlap_cfg=userbuffers_bf16_h100_h8192_tp4_mbs1_seqlen8192, # Predefined config + tp_comm_bootstrap_backend="nccl", # Communication backend +) +``` + +Requirements for TP communication overlap: +- `tensor_model_parallel_size >= 2` +- `sequence_parallel=True` +- Appropriate hardware configuration + +### Advanced Configuration + +For most use cases, setting `tp_comm_overlap=True` with `tp_comm_overlap_cfg=None` (the default) will automatically configure appropriate overlap settings. For advanced users requiring custom optimization, Megatron Bridge includes predefined configurations optimized for specific hardware and model combinations. These configurations are available in the `comm_overlap` module but require expert knowledge to use effectively. + +## Pipeline-parallel Communication Overlap + +Pipeline parallelism introduces P2P activation (gradient) sends and receives between pipeline-parallel (PP) GPUs. The PP communication frequency increases when increasing the virtual-pipeline-parallel size because the number of Transformer layers executed per micro-batch decreases. + +![Pipeline-parallel Communication Overlap](/assets/training/images/pp_comm_overlap.png) +*Figure: Pipeline-parallel communication overlap in 1F1B pipelining phase.* + +Megatron Bridge supports the overlap of PP communications with non-dependent computations in the 1F1B stage (the body of pipelining, where 1 forward and 1 backward micro-batch executions are interleaved). The PP communications in pipeline fill and flush stages are still exposed. + +### Configuration + +```python +comm_overlap_config = CommOverlapConfig( + tp_comm_overlap=False, + overlap_p2p_comm=True, # Enable PP communication overlap + batch_p2p_comm=False, # Use separate send/receive kernels +) +``` + +PP communication overlap settings: +- `overlap_p2p_comm`: Enables overlap of P2P communications (default: auto-configured) +- `batch_p2p_comm`: Uses batched vs separate kernels (default: auto-configured based on virtual PP) + +The overlap is automatically enabled when: +- `pipeline_model_parallel_size > 1` +- `virtual_pipeline_model_parallel_size > 1` (for optimal performance) + +## Context-parallel Communication Overlap + +Context parallelism partitions activations (gradients) on all layers in the sequence domain. This introduces all-gather and reduce-scatter of activations (gradients) in self-attention forward- and back-propagations. + +Megatron Bridge hides the context-parallel (CP) communications under the self-attention computation. Like the TP communication overlaps, the CP communications are chunked then pipeline-overlapped with the self-attention computation, where the all-gather and the reduce-scatter of activations (gradients) are replaced with P2P ring exchanges of data. + +### Automatic Configuration + +The CP communication overlap is automatically enabled when context parallelism is used (`context_parallel_size > 1`). No additional configuration is required as the overlap is built into the context parallelism implementation. + +## MoE Expert Parallel Communication Overlap + +For Mixture of Experts (MoE) models, Megatron Bridge supports overlapping expert parallel all-to-all communications with computation. + +### Configuration + +```python +comm_overlap_config = CommOverlapConfig( + tp_comm_overlap=False, + overlap_moe_expert_parallel_comm=True, # Enable MoE EP overlap + delay_wgrad_compute=True, # Advanced MoE optimization +) +``` + +Requirements for MoE expert parallel overlap: +- `expert_model_parallel_size > 1` +- `num_moe_experts > 1` +- `moe_token_dispatcher_type` in ["alltoall", "flex"] +- BF16 or FP16 precision +- PyTorch >= 2.6.0 +- Specific recomputation settings + +## Complete Configuration Example + +Here's a comprehensive example combining multiple communication overlap strategies: + +```python +from megatron.bridge.training.config import ConfigContainer, OptimizerConfig +from megatron.bridge.training.comm_overlap import ( + CommOverlapConfig, + userbuffers_bf16_h100_h8192_tp4_mbs1_seqlen8192 +) +from megatron.bridge.models import GPTModelProvider + +# Model configuration with parallelism +model_config = GPTModelProvider( + # Parallelism settings + tensor_model_parallel_size=4, + pipeline_model_parallel_size=2, + virtual_pipeline_model_parallel_size=2, + context_parallel_size=2, + sequence_parallel=True, + + # Model parameters + hidden_size=8192, + num_layers=32, + # ... other model parameters +) + +# Communication overlap configuration +comm_overlap_config = CommOverlapConfig( + # Tensor parallel overlap + tp_comm_overlap=True, + tp_comm_overlap_cfg=userbuffers_bf16_h100_h8192_tp4_mbs1_seqlen8192, + + # Pipeline parallel overlap + overlap_p2p_comm=True, + batch_p2p_comm=False, + + # Data parallel overlap + overlap_grad_reduce=True, + overlap_param_gather=True, + bucket_size=128 * 1024 * 1024, +) + +# Optimizer with distributed settings +optimizer_config = OptimizerConfig( + optimizer="adam", + lr=3e-4, + use_distributed_optimizer=True, +) + +# Complete configuration +config = ConfigContainer( + model=model_config, + comm_overlap=comm_overlap_config, + optimizer=optimizer_config, +) +``` + +## API Reference + +For detailed API documentation, see: +- [`bridge.training.comm_overlap.CommOverlapConfig`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.comm_overlap.CommOverlapConfig.html) - Main configuration class +- [`bridge.training.comm_overlap.TransformerLayerTPOverlapCfg`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.comm_overlap.TransformerLayerTPOverlapCfg.html) - Tensor parallel overlap configuration +- [`bridge.training.comm_overlap.BulkOverlapCfg`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.comm_overlap.BulkOverlapCfg.html) - Bulk overlap configuration +- [`bridge.training.comm_overlap.PipelineOverlapCfg`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.comm_overlap.PipelineOverlapCfg.html) - Pipeline overlap configuration +- [`bridge.training.comm_overlap.RingExchangeOverlapCfg`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.comm_overlap.RingExchangeOverlapCfg.html) - Ring exchange overlap configuration +megatron-core/developer-guide/latest/api-guide/tensor_parallel.html) - Underlying implementation details diff --git a/fern/v0.2.0/pages/training/config-container-overview.mdx b/fern/v0.2.0/pages/training/config-container-overview.mdx new file mode 100644 index 0000000000..1e76fcdfb3 --- /dev/null +++ b/fern/v0.2.0/pages/training/config-container-overview.mdx @@ -0,0 +1,199 @@ +--- +title: "Configuration Overview" +description: "" +--- + +The `ConfigContainer` is the central configuration object in Megatron Bridge that holds all settings for training. It acts as a single source of truth that brings together model architecture, training parameters, data loading, optimization, checkpointing, logging, and distributed training settings. + +## What is ConfigContainer + +`ConfigContainer` is a dataclass that holds all the configuration objects needed for training: + +```python +from megatron.bridge.training.config import ConfigContainer + +# ConfigContainer brings together all training configurations +config = ConfigContainer( + model=model_provider, # Model architecture and parallelism + train=training_config, # Training loop parameters + optimizer=optimizer_config, # Megatron Optimization settings + scheduler=scheduler_config, # Learning rate scheduling + dataset=dataset_config, # Data loading configuration + logger=logger_config, # Logging and monitoring + tokenizer=tokenizer_config, # Tokenization settings + checkpoint=checkpoint_config, # Checkpointing and resuming + dist=distributed_config, # Distributed training setup + ddp=ddp_config, # Megatron Distributed Data Parallel settings + # Optional configurations + peft=peft_config, # Parameter-efficient fine-tuning + profiling=profiling_config, # Performance profiling + mixed_precision=mp_config, # Mixed precision training + comm_overlap=comm_overlap_config, # Communication overlap settings + # ... and more +) +``` + +## Configuration Components + +| Component | Purpose | Required | Default | +|-----------|---------|----------|---------| +| `model` | Model architecture and parallelism strategy (GPT, T5, Mamba) | ✅ | - | +| `train` | Training loop parameters (batch sizes, iterations, validation) | ✅ | - | +| `optimizer` | Optimizer type and hyperparameters (from Megatron Core) | ✅ | - | +| `scheduler` | Learning rate and weight decay scheduling | ✅ | - | +| `dataset` | Data loading and preprocessing configuration | ✅ | - | +| `logger` | Logging, TensorBoard, and WandB configuration | ✅ | - | +| `tokenizer` | Tokenizer settings and vocabulary | ✅ | - | +| `checkpoint` | Checkpointing, saving, and loading | ✅ | - | +| `dist` | Distributed training initialization | | `DistributedInitConfig()` | +| `ddp` | Data parallel configuration (from Megatron Core) | | `DistributedDataParallelConfig()` | +| `rng` | Random number generation settings | | `RNGConfig()` | +| `rerun_state_machine` | Result validation and error injection | | `RerunStateMachineConfig()` | +| `mixed_precision` | Mixed precision training settings | | `None` | +| `comm_overlap` | Communication overlap optimizations | | `None` | +| `peft` | Parameter-efficient fine-tuning (LoRA, DoRA, etc.) | | `None` | +| `profiling` | Performance profiling with nsys or PyTorch profiler | | `None` | +| `ft` | Fault tolerance and automatic recovery | | `None` | +| `straggler` | GPU straggler detection | | `None` | +| `nvrx_straggler` | NVIDIA Resiliency Extension straggler detection | | `None` | +| `inprocess_restart` | In-process restart for fault tolerance | | `None` | + +## Design Philosophy + +### **Interoperability with External Config Systems** + +Megatron Bridge's Python configurations are designed to be amenable to other configuration systems you already use, such as: + +- Programmatic configuration: Direct Python object manipulation +- argparse: Command-line arguments can be easily mapped to dataclass fields +- File-based overrides: JSON, YAML, or other config files can override Python configs + +All of these approaches can be translated into Python dataclass instances. The framework provides utilities as a convenience for YAML-based overrides with OmegaConf, but the framework is not tied to any particular configuration system. + +```python +# All of these approaches work seamlessly: + +# 1. Direct Python configuration +config = ConfigContainer( + model=GPTModelProvider(num_layers=24, hidden_size=2048), + train=TrainingConfig(global_batch_size=256, train_iters=10000), + # ... other configs +) + +# 2. YAML-based serialization and deserialization (round-trip) +config.to_yaml("my_config.yaml") +config = ConfigContainer.from_yaml("my_config.yaml") # Load previously saved config + +# 3. Programmatic override after creation +config.train.global_batch_size = 512 # Override after instantiation +config.model.num_layers = 48 # Modify model architecture +``` + +### Centralized Configuration + +Megatron provides extensive flexibility through a rich set of configuration options. The `ConfigContainer` brings all these settings together in a single, organized object. This centralization makes configuration discoverable and maintainable - you have one place to understand and control all aspects of your training run. + +Unlike pure YAML-based configuration systems, `ConfigContainer` provides centralization with the full power of Python. You get the organizational benefits of a single configuration file combined with the programmatic flexibility of Python. + +The configuration system is built using nested dataclasses, providing: + +- **Modularity**: Each config component is independently defined and testable +- **Type safety**: Full static type checking +- **IDE support**: Autocomplete and type hints in development environments +- **Serialization**: Easy conversion to/from YAML, JSON, or other formats +- **Validation**: Built-in field validation + +```python +@dataclass +class ConfigContainer: + model: GPTModelProvider # Dataclass for model architecture + train: TrainingConfig # Dataclass for training parameters + optimizer: OptimizerConfig # Dataclass for optimization settings + # ... nested dataclasses for each concern +``` + +### Lazy Configuration and Deferred Validation + +For training workloads, configurations are lazy to support flexible user workflows: + +**Problem with Eager Validation:** +```python +# This would be problematic with eager validation: +config = TrainingConfig(train_iters=1000) +# __post_init__ calculates dependent values immediately + +config.train_iters = 5000 # User override +# Dependent values are now stale and incorrect! +``` + +**Solution with Lazy Finalization:** +```python +# Megatron Bridge approach - deferred validation +config = TrainingConfig(train_iters=1000) +config.train_iters = 5000 # User can safely override + +# Validation happens automatically right when training starts +pretrain(config, forward_step_func) # All dependent values calculated correctly +``` + +**Benefits:** +- Users can instantiate configs and subsequently override fields safely +- Dependent values are calculated correctly after all user modifications are applied +- Validation happens at the right time, right before training begins +- Flexible configuration workflows are supported + +### **Model Independence** + +Model configurations are designed to be independently usable outside the full training loop provided by thr framework: + +```python +# Models can be used standalone +model_provider = GPTModelProvider( + num_layers=24, + hidden_size=2048, + vocab_size=50000, # Must be explicitly set + seq_length=2048, # Must be explicitly set +) + +# This works independently of other configs +model_provider.finalize() +model = model_provider.provide() +``` + +**Trade-off**: The price for this flexibility is the need to explicitly set values like `seq_length` in multiple places during training. These settings are checked for consistency at the beginning of training. + +## Usage + +```python +# Create and configure +config = ConfigContainer( + model=GPTModelProvider(num_layers=24, seq_length=2048), + train=TrainingConfig(train_iters=1000), + dataset=GPTDatasetConfig(seq_length=2048), # Must match model seq_length + # ... other required configs +) + +# Modify as needed +config.train.train_iters = 5000 +config.model.hidden_size = 4096 + +# Start training - validation happens automatically +pretrain(config, forward_step_func) +``` + +## Configuration Export and Import + +### Export to YAML +```python +# Print YAML configuration to console +config.print_yaml() + +# Save to file +config.to_yaml("config.yaml") +``` + +### Load from YAML +```python +# Load configuration from YAML file +config = ConfigContainer.from_yaml("config.yaml") +``` diff --git a/fern/v0.2.0/pages/training/cpu-offloading.mdx b/fern/v0.2.0/pages/training/cpu-offloading.mdx new file mode 100644 index 0000000000..7f9b42ede1 --- /dev/null +++ b/fern/v0.2.0/pages/training/cpu-offloading.mdx @@ -0,0 +1,79 @@ +--- +title: "CPU Offloading" +description: "" +--- + +## Overview + +CPU Offloading in Megatron Bridge is a feature that reduces the peak memory usage of the GPU by offloading activations and inactive weights to CPU storage. Megatron Bridge supports offloading at the transformer layer level, allowing users to specify the number of transformer layers in their language model that require CPU offloading. During the forward pass, Megatron Bridge offloads activations at the optimal time and reloads them as needed during the backward pass. + +## Features + +- Supports training models with long sequence lengths by managing activation memory efficiently +- Enables high batch sizes per GPU by offloading activation memory +- Overlaps computation with data transfers (Host2Device and Device2Host) during offloading and reloading + +## Configuration + +CPU offloading is configured through the model provider parameters: + +```python +from megatron.bridge.models import GPTModelProvider + +# Basic CPU offloading configuration +model_config = GPTModelProvider( + # Model architecture + hidden_size=4096, + num_layers=32, + + # CPU offloading settings + cpu_offloading=True, # Enable CPU offloading + cpu_offloading_num_layers=16, # Number of layers to offload (0 to num_layers-1) + cpu_offloading_activations=True, # Offload activations + cpu_offloading_weights=True, # Offload weights + + # ... other model parameters +) +``` + +### Configuration Parameters + +- **`cpu_offloading`**: Set to `True` to enable CPU offloading +- **`cpu_offloading_num_layers`**: Number of transformer layers to offload (value between 0 and total number of layers minus one) +- **`cpu_offloading_activations`**: Whether to offload activations to CPU memory (default: `True`) +- **`cpu_offloading_weights`**: Whether to offload inactive weights to CPU memory (default: `False`) +- **`cpu_offloading_double_buffering`**: Enable double buffering across layers while reloading activations from CPU (default: `False`) + +### Offloading Strategies + +You can configure different combinations of offloading based on your memory requirements: + +#### Activations Only +```python +model_config = GPTModelProvider( + cpu_offloading=True, + cpu_offloading_num_layers=8, + cpu_offloading_activations=True, # Offload activations + cpu_offloading_weights=False, # Keep weights on GPU +) +``` + +#### Weights Only +```python +model_config = GPTModelProvider( + cpu_offloading=True, + cpu_offloading_num_layers=8, + cpu_offloading_activations=False, # Keep activations on GPU + cpu_offloading_weights=True, # Offload weights +) +``` + +#### Both Activations and Weights +```python +model_config = GPTModelProvider( + cpu_offloading=True, + cpu_offloading_num_layers=8, + cpu_offloading_activations=True, # Offload activations + cpu_offloading_weights=True, # Offload weights +) +``` diff --git a/fern/v0.2.0/pages/training/distillation.mdx b/fern/v0.2.0/pages/training/distillation.mdx new file mode 100644 index 0000000000..331bbed4a9 --- /dev/null +++ b/fern/v0.2.0/pages/training/distillation.mdx @@ -0,0 +1,124 @@ +--- +title: "Knowledge Distillation" +description: "" +--- + +Megatron Bridge provides a streamlined setup for Knowledge Distillation (KD) training, making it easy to enable and integrate into your workflow. This section explains how to use this feature effectively. + +Knowledge Distillation is a technique where a pre-trained model (the "teacher") transfers its learned knowledge to a second model (the "student"), which is typically smaller and faster. This process helps the student model learn more efficiently by mimicking the behavior of the teacher. KD offers two key advantages over traditional training: faster convergence and higher final accuracy. + +In Megatron Bridge, KD is enabled by NVIDIA Model Optimizer (ModelOpt) — a library to optimize deep-learning models for inference on GPUs. + +## Knowledge Distillation Process + +The KD process involves these steps: + +1. **Loads Checkpoints**: Loads both the student and teacher model checkpoints. +2. **Replaces Loss Function**: Replaces the standard loss function with the KL-Divergence between the output logits (and potentially additional losses between pairs of intermediate model states). +3. **Trains Models**: Runs forward passes on both models, but executes the backward pass only on the student model. +4. **Saves Checkpoints**: Saves only the student model checkpoint, allowing it to be used later in the same manner as before. + +## Limitations + +* Only GPT-based checkpoints are currently supported. +* Student and teacher models must support the same parallelism strategy. +* If Pipeline Parallelism is enabled, intermediate-state based KD losses are only supported on the final pipeline stage. + +## Configuration + +### Knowledge Distillation Config + +You can configure the KD process via the `ModelOptDistillConfig` class or a YAML file. The configuration includes: + +* `logit_layers`: The layer names of student and teacher model logit layers. These names correspond to the PyTorch submodule attributes of the Megatron Core model. (For GPT-based models, this is `"output_layer"`). Default: `["output_layer", "output_layer"]` +* `intermediate_layer_pairs`: A list of pairs of intermediate layer names. These pairs will by default have a Cosine-Similarity loss between them, and if tensor-parallelism is enabled, these layers must have sequence parallel outputs (i.e. LayerNorms), as Cosine loss cannot have a split hidden dimension. Default: `[["decoder.final_layernorm", "decoder.final_layernorm"]]` +* `skip_lm_loss`: Whether to skip the default language modeling (LM) loss. If `false`, it will be added to the distillation loss. (Note it consumes more memory). Default: `true` +* `kd_loss_scale`: Relative scale factor for the distillation loss. The cumulative logits-and-intermediate loss gets scaled to `kd_loss_scale` times the magnitude of the LM loss. Not used if `skip_lm_loss` is `true`. Default: `1.0` +* `logit_kl_temperature`: Temperature variable for KL Divergence loss calculation. Default: `1.0` + +Example YAML configuration: + +```yaml +logit_layers: ["output_layer", "output_layer"] +intermediate_layer_pairs: + - ["decoder.final_layernorm", "decoder.final_layernorm"] +logit_kl_temperature: 2.0 +``` + +## Usage + +### Basic Usage with Default Configuration + +The simplest way to run knowledge distillation is to use or adapt one of the provided recipe scripts. Here's an example for distilling Llama3.2-3B into Llama3.2-1B: + +```bash +torchrun --nproc_per_node=1 examples/distillation/llama/distill_llama32_3b-1b.py +``` + +### Using a Custom YAML Config File + +You can provide a custom YAML configuration file to override default settings: + +```bash +torchrun --nproc_per_node=1 examples/distillation/llama/distill_llama32_3b-1b.py \ + --config-file my_custom_config.yaml +``` + +### Using CLI Overrides + +Megatron Bridge supports Hydra-style CLI overrides for flexible configuration: + +```bash +torchrun --nproc_per_node=2 examples/distillation/llama/distill_llama32_3b-1b.py \ + model.tensor_model_parallel_size=2 \ + model.teacher.tensor_model_parallel_size=2 +``` + +### Combining YAML and CLI Overrides + +CLI overrides take precedence over YAML configuration: + +```bash +torchrun --nproc_per_node=2 examples/distillation/llama/distill_llama32_3b-1b.py \ + --config-file conf/my_config.yaml \ + train.global_batch_size=512 +``` + +## Model Support + +Currently, distillation is supported for GPT and Mamba-based models + +To enable distillation for a model: + +1. Set the `teacher` attribute to the teacher model configuration +2. Configure `kd_config` with desired distillation settings (else uses default) +3. Use `convert_to_distillation_provider()` to convert your existing model provider + +## Checkpointing + +During distillation training: + +* Only the **student model** checkpoints are saved +* Teacher model remains frozen and is not modified +* Checkpoints can be used for inference or further training like any standard checkpoint + +## Best Practices + +1. **Match Parallelism**: Ensure student and teacher use compatible parallelism configurations +2. **Monitor Loss**: Track both distillation loss and (if enabled) language modeling loss +3. **Batch Size**: Use larger batch sizes for better stability during distillation +4. **Learning Rate**: Start with a smaller LR than pretraining +5. **Data Quality**: Use high-quality, diverse training data for best distillation results + +## Troubleshooting + +### Out of Memory Errors + +* Reduce `train.micro_batch_size` +* Increase parallelism sizes +* Set `model.kd_config.skip_lm_loss = True` to save memory + +## References + +For more information on the underlying implementation, see: +* [NVIDIA Model Optimizer](https://github.com/NVIDIA/Model-Optimizer) diff --git a/fern/v0.2.0/pages/training/entry-points.mdx b/fern/v0.2.0/pages/training/entry-points.mdx new file mode 100644 index 0000000000..3752023b60 --- /dev/null +++ b/fern/v0.2.0/pages/training/entry-points.mdx @@ -0,0 +1,99 @@ +--- +title: "Training Entry Points" +description: "" +--- + +Megatron Bridge provides unified training entry points for pretraining, Supervised Fine-Tuning (SFT), and Parameter-Efficient Fine-Tuning (PEFT). All training modes share the same underlying training loop architecture, differing primarily in their data handling and model configuration. + +## Main Entry Points + +The [`bridge.training.pretrain.pretrain`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.pretrain.pretrain.html) and [`bridge.training.finetune.finetune`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.finetune.finetune.html) functions are the primary entry points for pretraining models—either from scratch or through fine-tuning. Each function accepts a [`bridge.training.config.ConfigContainer`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.ConfigContainer.html) along with a `forward_step_func` that defines how the training loop should be run. + +## Forward Step Function + +The `forward_step_func` defines how each training step is executed. It should follow this signature: + +```python +def forward_step_func( + global_state: GlobalState, + data_iterator: Iterable, + model: MegatronModule, + return_schedule_plan: bool = False, +) -> tuple[Any, Callable]: + """Forward step function. + + Args: + global_state: Training state object containing configuration and utilities + data_iterator: Iterator over training/evaluation data + model: The model to perform forward step on + return_schedule_plan: Whether to return schedule plan (for MoE overlap) + + Returns: + tuple containing: + - output: Forward pass output (tensor or collection of tensors) + - loss_func: Function to compute loss from the output + """ +``` + +### Responsibilities + +The forward step function has three main responsibilities: + +1. **Get a Batch**: Retrieve and process the next batch from the data iterator. +2. **Run Forward Pass**: Execute the model's forward pass on the batch. +3. **Return Loss Function**: Provide a function to compute loss from the output. + +### State Access + +Megatron Bridge automatically provides the [`bridge.training.state.GlobalState`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.state.GlobalState.html) object containing: +- **Configuration**: Complete training configuration (`global_state.cfg`). +- **Timers**: Performance monitoring utilities (`global_state.timers`). +- **Training Progress**: Current step, consumed samples (`global_state.train_state`). +- **Loggers**: TensorBoard and WandB loggers for metrics tracking. + +All configuration and state information are accessible through the injected `state` object. + +For complete implementation examples, see [`bridge.training.gpt_step.forward_step`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.gpt_step.forward_step.html). + +## Loss Calculation and Reduction + +The loss function returned by the forward step can follow different patterns based on your needs: + +### Loss Function Patterns + +1. **Standard Pattern**: Return `(loss, metadata_dict)` + - The loss is automatically averaged across microbatches + - Metadata dict contains named loss components for logging + - Most common pattern for standard training + +2. **Token-aware Pattern**: Return `(loss, num_tokens, metadata_dict)` + - Loss is averaged across both microbatches and tokens + - Useful when you want per-token loss averaging + - Recommended for variable-length sequences + +3. **Inference Pattern**: Return arbitrary data structures + - Used with `collect_non_loss_data=True` and `forward_only=True` + - Suitable for inference, evaluation metrics, or custom data collection + - No automatic loss processing applied + +### Automatic Loss Processing + +The training loop automatically handles: +- **Microbatch Reduction**: Aggregates losses across all microbatches in the global batch. +- **Distributed Reduction**: Performs all-reduce operations across data parallel ranks. +- **Pipeline Coordination**: Only the last pipeline stage computes and reduces losses. +- **Logging Integration**: Automatically logs loss components to TensorBoard/WandB. + +For implementation details, see [`bridge.training.train.train_step`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.train.train_step.html) and [`bridge.training.losses.masked_token_loss`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.losses.masked_token_loss.html), as an example. + +## Customization + +### When to Customize + +You can customize the forward step function when you need: + +- **Custom Loss Functions**: Beyond standard language modeling loss (e.g., adding regularization, multi-objective training). +- **Multi-task Learning**: Training models on multiple tasks simultaneously with different loss components. +- **Custom Data Processing**: Specialized batch preprocessing for domain-specific data formats. +- **Additional Metrics**: Computing extra evaluation metrics during training. +- **Model-specific Logic**: Special handling for custom model architectures or training procedures. diff --git a/fern/v0.2.0/pages/training/logging.mdx b/fern/v0.2.0/pages/training/logging.mdx new file mode 100644 index 0000000000..4d633e9764 --- /dev/null +++ b/fern/v0.2.0/pages/training/logging.mdx @@ -0,0 +1,376 @@ +--- +title: "Logging and Monitoring" +description: "" +--- + +This guide describes how to configure logging in Megatron Bridge. It introduces the high-level `LoggerConfig`, explains experiment logging to TensorBoard and Weights & Biases (W&B), and documents console logging behavior. + +## LoggerConfig Overview + +[`bridge.training.config.LoggerConfig`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.LoggerConfig.html) is the dataclass that encapsulates logging‑related settings for training. It resides inside the overall [`bridge.training.config.ConfigContainer`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.ConfigContainer.html), which represents the complete configuration for a training run. + +### Timer Configuration Options + +Use the following options to control which timing metrics are collected during training and how they are aggregated and logged. + +#### `timing_log_level` +Controls which timers are recorded during execution: + +- **Level 0**: Logs only the overall iteration time. +- **Level 1**: Includes once-per-iteration operations, such as gradient all-reduce. +- **Level 2**: Captures frequently executed operations, providing more detailed insights but with increased overhead. + +#### `timing_log_option` +Specifies how timer values are aggregated across ranks. Valid options: + +- `"max"`: Logs the maximum value across ranks. +- `"minmax"`: Logs both minimum and maximum values. +- `"all"`: Logs all values from all ranks. + +#### `log_timers_to_tensorboard` +When enabled, the framework records timer metrics to supported backends such as TensorBoard. + +### Diagnostic Options + +The framework provides several optional toggles for enhanced monitoring and diagnostics: + +- **Loss Scale**: Enables dynamic loss scaling for mixed-precision training. +- **Validation Perplexity**: Tracks model perplexity during validation. +- **CUDA Memory Statistics**: Reports detailed GPU memory usage. +- **World Size**: Displays the total number of distributed ranks. + +### Logging Options + +Use the following options to enable additional diagnostics and performance monitoring during training. + +- **`log_params_norm`**: Computes and logs the L2 norm of model parameters. If available, it also logs the gradient norm. +- **`log_energy`**: Activates the energy monitor, which records per-GPU energy consumption and instantaneous power usage. +- **`log_memory`**: Logs the memory usage of the model from `torch.cuda.memory_stats()`. +- **`log_throughput_to_tensorboard`**: Calculates the training throughput and utilization. +- **`log_runtime_to_tensorboard`**: Estimates total time remaining until the end of the training. +- **`log_l2_norm_grad_to_tensorboard`**: Computes and logs the L2 norm of gradients for each model layer. + +## Experiment Logging +Both TensorBoard and W&B are supported for metric logging. When using W&B, it’s recommended to also enable TensorBoard to ensure that all scalar metrics are consistently logged across backends. + +### TensorBoard + + +#### What Gets Logged + +TensorBoard captures a range of training and system metrics, including: + +- **Learning rate**, including decoupled LR when applicable +- **Per-loss scalars** for detailed breakdowns +- **Batch size** and **loss scale** +- **CUDA memory usage** and **world size** (if enabled) +- **Validation loss**, with optional **perplexity** +- **Timers**, when timing is enabled +- **Energy consumption** and **instantaneous power**, if energy logging is active + +#### Enable TensorBoard Logging + 1) Install TensorBoard (if not already available): + ```bash + pip install tensorboard + ``` + 2) **Configure logging** in your training setup. In these examples, `cfg` refers to a `ConfigContainer` instance (such as one produced by a recipe), which contains a `logger` attribute representing the `LoggerConfig`: + + ```python + from megatron.bridge.training.config import LoggerConfig + + cfg.logger = LoggerConfig( + tensorboard_dir="./runs/tensorboard", + tensorboard_log_interval=10, + log_timers_to_tensorboard=True, # optional + log_memory_to_tensorboard=False, # optional + ) + ``` + + + The writer is created lazily on the last rank when `tensorboard_dir` is set. + + +#### Set the Output Directory + +TensorBoard event files are saved to the directory specified by `tensorboard_dir`. + +**Example with additional metrics enabled:** +```python +cfg.logger.tensorboard_dir = "./logs/tb" +cfg.logger.tensorboard_log_interval = 5 +cfg.logger.log_loss_scale_to_tensorboard = True +cfg.logger.log_validation_ppl_to_tensorboard = True +cfg.logger.log_world_size_to_tensorboard = True +cfg.logger.log_timers_to_tensorboard = True +``` + +### Weights & Biases (W&B) + + +#### What Gets Logged + +When enabled, W&B automatically mirrors the scalar metrics logged to TensorBoard. +In addition, the full run configuration is synced at initialization, allowing for reproducibility and experiment tracking. + +#### Enable W&B Logging + + 1) Install W&B (if not already available): + ```bash + pip install wandb + ``` + 2) Authenticate with W&B using one of the following methods: + - Set `WANDB_API_KEY` in the environment before the run, or + - Run `wandb login` once on the machine. + 2) **Configure logging** in your training setup. In these examples, `cfg` refers to a `ConfigContainer` instance (such as one produced by a recipe), which contains a `logger` attribute representing the `LoggerConfig`: + + ```python + from megatron.bridge.training.config import LoggerConfig + + cfg.logger = LoggerConfig( + tensorboard_dir="./runs/tensorboard", # recommended: enables shared logging gate + wandb_project="my_project", + wandb_exp_name="my_experiment", + wandb_entity="my_team", # optional + wandb_save_dir="./runs/wandb", # optional + ) + ``` + + +W&B is initialized lazily on the last rank when `wandb_project` is set and `wandb_exp_name` is non-empty. + + +#### W&B Configuration with NeMo Run Launching + +For users launching training scripts with NeMo Run, W&B can be optionally configured using the [`bridge.recipes.run_plugins.WandbPlugin`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.recipes.run_plugins.WandbPlugin.html). + +The plugin automatically forwards the `WANDB_API_KEY` and by default injects CLI overrides for the following logger parameters: + +- `logger.wandb_project` +- `logger.wandb_entity` +- `logger.wandb_exp_name` +- `logger.wandb_save_dir` + +This allows seamless integration of W&B logging into your training workflow without manual configuration. + +### MLFlow + +Megatron Bridge can log metrics and artifacts to MLFlow, following the same pattern as the W&B integration. + +#### What Gets Logged + +When enabled, MLFlow receives: + +- Training configuration as run parameters +- Scalar metrics (losses, learning rate, batch size, throughput, timers, memory, runtime, norms, energy, etc.) +- Checkpoint artifacts saved under an experiment-specific artifact path per iteration + +#### Enable MLFlow Logging + + 1) Install MLFlow (installed by default with Megatron Bridge): + + ```bash + pip install mlflow / uv add mlflow + ``` + + 2) Configure the tracking server (Optional): + - Either set `MLFLOW_TRACKING_URI` in the environment, or + - Pass an explicit `mlflow_tracking_uri` in the logger config. + + 3) Configure logging in your training setup. + + ```python + from megatron.bridge.training.config import LoggerConfig + + cfg.logger = LoggerConfig( + tensorboard_dir="./runs/tensorboard", + mlflow_experiment="my_megatron_experiment", + mlflow_run_name="llama32_1b_pretrain_run", + mlflow_tracking_uri="http://mlflow:5000", # optional + mlflow_tags={ # optional + "project": "llama32", + "phase": "pretrain", + }, + ) + ``` + +#### Progress Log + +When `logger.log_progress` is enabled, the framework generates a `progress.txt` file in the checkpoint save directory. + +This file includes: +- **Job-level metadata**, such as timestamp and GPU count +- **Periodic progress entries** throughout training + +At each checkpoint boundary, the log is updated with: +- **Job throughput** (TFLOP/s/GPU) +- **Cumulative throughput** +- **Total floating-point operations** +- **Tokens processed** + +This provides a lightweight, text-based audit trail of training progress, useful for tracking performance across restarts. + +## Tensor Inspection + +Megatron Bridge integrates with TransformerEngine's tensor inspection features via NVIDIA DLFW Inspect. This integration, controlled by [`bridge.training.config.TensorInspectConfig`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.TensorInspectConfig.html), enables advanced debugging and analysis of tensor statistics during training. When enabled, the framework handles initialization, step tracking, and cleanup automatically. + + +**Current limitations:** Tensor inspection is currently supported only for linear modules in TransformerEngine (e.g., `fc1`, `fc2`, `layernorm_linear`). Operations like attention are not supported. + + + +This section covers Megatron Bridge configuration. For comprehensive documentation on features, configuration syntax, and advanced usage, see: + +- [TransformerEngine Debug Documentation](https://github.com/NVIDIA/TransformerEngine/tree/af2a0c16ec11363c0af84690cd877a59f898820e/docs/debug) +- [NVIDIA DLFW Inspect Documentation](https://github.com/NVIDIA/nvidia-dlfw-inspect/tree/4118044cc84f0183714a2ab1bc215fa49f9aaa82/docs) + + +### Installation + +Install NVIDIA DLFW Inspect if not already available: +```bash +pip install nvdlfw-inspect +``` + +### Available Features + +TransformerEngine provides the following debug features: + +- **LogTensorStats** – Logs high-precision tensor statistics: `min`, `max`, `mean`, `std`, `l1_norm`, `l2_norm`, `cur_amax`, `dynamic_range`. +- **LogFp8TensorStats** – Logs quantized tensor statistics for FP8 recipes: `underflows%`, `scale_inv_min`, `scale_inv_max`, `mse`. Supports simulating alternative recipes (e.g., tracking `mxfp8_underflows%` during per-tensor current-scaling training) +- **DisableFP8GEMM** – Runs specific GEMM operations in high precision +- **DisableFP8Layer** – Disables FP8 for entire layers +- **PerTensorScaling** – Enables per-tensor current scaling for specific tensors +- **FakeQuant** – Experimental quantization testing + +See [TransformerEngine debug features](https://github.com/NVIDIA/TransformerEngine/tree/af2a0c16ec11363c0af84690cd877a59f898820e/transformer_engine/debug/features) for complete parameter lists and usage details. + +### Configuration + +Configure tensor inspection using [`bridge.training.config.TensorInspectConfig`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.TensorInspectConfig.html) with either a YAML file or inline dictionary. + +#### YAML Configuration + +```yaml +tensor_inspect: + enabled: true + features: ./conf/fp8_tensor_stats.yaml + log_dir: ./logs/tensor_inspect +``` + +**Example feature configuration file:** + +```yaml +fp8_tensor_stats: + enabled: true + layers: + layer_name_regex_pattern: ".*(fc2)" + transformer_engine: + LogFp8TensorStats: + enabled: true + tensors: [weight,activation,gradient] + stats: ["underflows%", "mse"] + freq: 5 + start_step: 0 + end_step: 100 +``` + +#### Python Configuration + +```python +from bridge.training.config import TensorInspectConfig + +# Option 1: inline python dict +cfg.tensor_inspect = TensorInspectConfig( + enabled=True, + features={ + "fp8_gradient_stats": { + "enabled": True, + "layers": {"layer_name_regex_pattern": ".*(fc1|fc2)"}, + "transformer_engine": { + "LogFp8TensorStats": { + "enabled": True, + "tensors": ["weight","activation","gradient"], + "stats": ["underflows%", "mse"], + "freq": 5, + "start_step": 0, + "end_step": 100, + }, + }, + } + }, + log_dir="./logs/tensor_inspect", +) + +# Option 2: reference external YAML +cfg.tensor_inspect = TensorInspectConfig( + enabled=True, + features="./conf/fp8_inspect.yaml", + log_dir="./logs/tensor_inspect", +) + +``` + +#### Layer Selection + +Features apply to linear modules matched by selectors in the `layers` section: + +- `layer_name_regex_pattern: .*` – All supported linear layers +- `layer_name_regex_pattern: .*layers\.(0|1|2).*(fc1|fc2|layernorm_linear)` – Linear modules in first three transformer layers +- `layer_name_regex_pattern: .*(fc1|fc2)` – MLP projections only +- `layer_types: [layernorm_linear, fc1]` – String matching (alternative to regex) + +Tensor-level selectors (`tensors`, `tensors_struct`) control which tensor roles are logged: `activation`, `gradient`, `weight`, `output`, `wgrad`, `dgrad`. + +### Output and Monitoring + +Tensor statistics are written to `tensor_inspect.log_dir` and forwarded to TensorBoard/W&B when enabled. + +**Log locations:** +- Text logs: `/nvdlfw_inspect_statistics_logs/` +- TensorBoard +- W&B + +### Performance Considerations + +- Use `freq > 1` to reduce overhead. Statistics collection is expensive for large models. +- Narrow layer selection with specific regex patterns rather than `.*` + +## Console Logging + +Megatron Bridge uses the standard Python logging subsystem for console output. + +### Configure Console Logging + +To control console logging behavior, use the following configuration options: + +- `logging_level` sets the default verbosity level. It can be overridden via the `MEGATRON_BRIDGE_LOGGING_LEVEL` environment variable. +- `filter_warnings` suppresses messages at the WARNING level. +- `modules_to_filter` specifies logger name prefixes to exclude from output. +- `set_level_for_all_loggers` determines whether the logging level is applied to all loggers or only a subset, depending on the current implementation. + +### Monitor Logging Cadence and Content + +To monitor training progress at regular intervals, the framework prints a summary line every `log_interval` iterations. + +Each summary includes: +- **Timestamp** +- **Iteration counters** +- **Consumed and skipped samples** +- **Iteration time (ms)** +- **Learning rates** +- **Global batch size** +- **Per-loss averages** +- **Loss scale** + +When enabled, additional metrics are printed: +- **Gradient norm** +- **Zeros in gradients** +- **Parameter norm** +- **Energy and power per GPU** + +Straggler timing reports follow the same `log_interval` cadence, helping identify performance bottlenecks across ranks. + +### Minimize Timing Overhead + +To reduce performance impact, set `timing_log_level` to `0`. +Increase to `1` or `2` only when more detailed timing metrics are required, as higher levels introduce additional logging overhead. diff --git a/fern/v0.2.0/pages/training/mixed-precision.mdx b/fern/v0.2.0/pages/training/mixed-precision.mdx new file mode 100644 index 0000000000..ce629b8b2a --- /dev/null +++ b/fern/v0.2.0/pages/training/mixed-precision.mdx @@ -0,0 +1,190 @@ +--- +title: "Mixed Precision Training" +description: "" +--- + +Mixed precision training significantly enhances computational efficiency by conducting operations in low-precision format, while selectively maintaining minimal data in single-precision to preserve critical information throughout key areas of the network. Megatron Bridge supports FP16, BF16, and FP8 via Transformer Engine (TE) across most models through the [`bridge.training.mixed_precision.MixedPrecisionConfig`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.mixed_precision.MixedPrecisionConfig.html) configuration. + +## Configuration Overview + +Mixed precision is configured in Megatron Bridge through the `mixed_precision` field in [`bridge.training.config.ConfigContainer`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.ConfigContainer.html), which accepts either: +- A string name referencing a predefined recipe (e.g., `"bf16_mixed"`) +- A [`bridge.training.mixed_precision.MixedPrecisionConfig`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.mixed_precision.MixedPrecisionConfig.html) object for custom configurations + +The mixed precision configuration automatically updates the model, optimizer, and distributed data parallel settings with the appropriate precision parameters. + +## Half-Precision Training + +Megatron Bridge supports half-precision FP16 and BF16 computation training via Megatron Core and the distributed optimizer. This training recipe uses half-precision in all layer computation while keeping the model states (optimizer states and master parameters) in single-precision. To avoid repeated data type casting at each layer computation, Megatron Core keeps a separate copy of half-precision parameters that is updated after each optimizer step. + +### Using Predefined Recipes + +The simplest way to enable mixed precision is using predefined recipe names: + +```python +from megatron.bridge.training.config import ConfigContainer + +# Configure with BF16 mixed precision +config = ConfigContainer( + mixed_precision="bf16_mixed", + # ... other config parameters +) + +# Configure with FP16 mixed precision +config = ConfigContainer( + mixed_precision="fp16_mixed", + # ... other config parameters +) +``` + +### Custom Mixed Precision Configuration + +For more control, create a custom [`bridge.training.mixed_precision.MixedPrecisionConfig`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.mixed_precision.MixedPrecisionConfig.html): + +```python +from megatron.bridge.training.config import ConfigContainer +from megatron.bridge.training.mixed_precision import MixedPrecisionConfig +import torch + +# Custom BF16 configuration +bf16_config = MixedPrecisionConfig( + bf16=True, + params_dtype=torch.bfloat16, + pipeline_dtype=torch.bfloat16, + autocast_enabled=False, + grad_reduce_in_fp32=True, +) + +config = ConfigContainer( + mixed_precision=bf16_config, + # ... other config parameters +) +``` + +## FP8 Training + +NVIDIA H100 GPU introduced support for a new datatype, FP8 (8-bit floating point), enabling higher throughput of matrix multiplies and convolutions. Megatron Bridge uses the NVIDIA TransformerEngine (TE) to leverage speedups from FP8. For a more detailed overview, refer to the [TE documentation](https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/index.html), specifically the FP8 format and recipe. + +### FP8 Configuration Parameters + +The [`bridge.training.mixed_precision.MixedPrecisionConfig`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.mixed_precision.MixedPrecisionConfig.html) provides several FP8-specific parameters: + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `fp8` | `Optional[str]` | `None` | FP8 format: `"hybrid"` (E4M3 for activations/weights, E5M2 for gradients) or `"e4m3"` | +| `fp8_recipe` | `str` | `"tensorwise"` | FP8 recipe type: `"tensorwise"`, `"delayed"`, `"blockwise"`, `"mxfp8"` (Blackwell only) | +| `first_last_layers_bf16` | `bool` | `False` | If True, retains first and last N TransformerBlocks in BF16 as opposed to FP8 | +| `num_layers_at_start_in_bf16` | `int` | `0` | Number of layers at the start of the model to keep in BF16 precision when `first_last_layers_bf16` is True | +| `num_layers_at_end_in_bf16` | `int` | `0` | Number of layers at the end of the model to keep in BF16 precision when `first_last_layers_bf16` is True | +| `fp8_margin` | `int` | `0` | Scaling factor shift by $2^{margin}$ | +| `fp8_amax_history_len` | `int` | `1` | Window size for amax history storage | +| `fp8_amax_compute_algo` | `str` | `"most_recent"` | Amax selection algorithm: `"max"` or `"most_recent"` | +| `fp8_param` | `Optional[bool]` | `None` | Store module-level parameters in FP8 | +| `fp8_param_gather` | `bool` | `False` | Enable FP8 parameter gathering | + +### FP8 Recipe Examples + +Use any of the predefined FP8 recipe names with the `mixed_precision` parameter: + +```python +# Example: BF16 with FP8 current scaling +config = ConfigContainer( + mixed_precision="bf16_with_fp8_current_scaling_mixed", + # ... other config parameters +) +``` + +## Available Mixed Precision Recipes + +Megatron Bridge provides numerous predefined mixed precision recipes for different use cases. You can use the [`megatron.bridge.training.mixed_precision.get_mixed_precision_config`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/megatron.bridge.training.mixed_precision.get_mixed_precision_config.html) utility function to convert from a string shortname to a class instance. For the complete list of available recipes and their specific configurations, see the [`megatron.bridge.training.mixed_precision`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/megatron.bridge.training.mixed_precision.html) module. + +### Custom FP8 Configuration + +For advanced use cases, create a custom FP8 configuration: + +```python +from megatron.bridge.training.mixed_precision import MixedPrecisionConfig +import torch + +# Custom FP8 configuration +fp8_config = MixedPrecisionConfig( + bf16=True, + params_dtype=torch.bfloat16, + pipeline_dtype=torch.bfloat16, + fp8="hybrid", + fp8_recipe="tensorwise", + fp8_margin=0, + fp8_amax_history_len=1024, + fp8_amax_compute_algo="max", + fp8_param_gather=True, +) + +config = ConfigContainer( + mixed_precision=fp8_config, + # ... other config parameters +) +``` + +### Registering Custom Mixed Precision Recipes + +You can also register your own custom mixed precision configurations to work with the shortname system. Use the [`megatron.bridge.training.mixed_precision.register`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/megatron.bridge.training.mixed_precision.register.html) decorator on a function that returns a `MixedPrecisionConfig` object: + +```python +from megatron.bridge.training.mixed_precision import register, MixedPrecisionConfig + +@register +def my_custom_fp8_recipe() -> MixedPrecisionConfig: + """Custom FP8 recipe with specific settings for my use case.""" + return MixedPrecisionConfig( + bf16=True, + fp8="hybrid", + fp8_recipe="tensorwise", + fp8_param_gather=True, + # ... other custom settings + ) + +# Now you can use it with the utility function +config = get_mixed_precision_config("my_custom_fp8_recipe") +``` + +Common recipe categories include: +- **Half-precision recipes**: Basic BF16 and FP16 mixed precision +- **FP8 recipes**: Various FP8 scaling strategies (delayed, current, subchannel) +- **Architecture-specific recipes**: Optimized for specific GPU architectures (Hopper, Blackwell) +- **Model-specific recipes**: Tuned for particular model families + +## Configuration Synchronization + +When a mixed precision configuration is provided, it automatically synchronizes precision-related settings across the model, optimizer, and distributed data parallel (DDP) configurations. This ensures consistent precision behavior throughout the training pipeline. + +**Important**: Mixed precision settings will override any conflicting precision parameters that may have been set directly on the model, optimizer, or DDP configurations. The mixed precision configuration acts as the authoritative source for all precision-related parameters. + +For example, if you specify both: +```python +# This will be overridden +model_config.bf16 = False +optimizer_config.bf16 = False + +config = ConfigContainer( + model=model_config, + optimizer=optimizer_config, + mixed_precision="bf16_mixed", # This takes precedence during training + # ... other configs +) +``` + +The mixed precision configuration will set `bf16=True` on both the model and optimizer configs, overriding the explicitly set `False` values. This synchronization prevents configuration mismatches that could lead to training issues. + +## Performance Considerations + +- **FP8 recipes are experimental** and convergence has not been fully validated for all models +- **BF16** is generally recommended over FP16 for better numerical stability +- **FP8** provides the best performance on H100 GPUs but requires careful tuning +- **MXFP8** recipes are only supported on Blackwell architecture GPUs +- **Blockwise scaling** recipes are optimized for Hopper architecture GPUs + +## Resources + +- [Transformer Engine Documentation](https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/index.html) +- [Intro to FP8, floating point formats, and mixed precision training](https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/examples/fp8_primer.html#Introduction-to-FP8) +- [Performance optimizations](https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/examples/advanced_optimizations.html) that are natively supported in Megatron Bridge by enabling FP8 training with TE diff --git a/fern/v0.2.0/pages/training/multi-token-prediction.mdx b/fern/v0.2.0/pages/training/multi-token-prediction.mdx new file mode 100644 index 0000000000..d9e31de5a2 --- /dev/null +++ b/fern/v0.2.0/pages/training/multi-token-prediction.mdx @@ -0,0 +1,274 @@ +--- +title: "Multi-Token Prediction (MTP)" +description: "" +--- + +## Overview + +Multi-Token Prediction (MTP) is an advanced training technique introduced in the [DeepSeek-V3 Technical Report](https://arxiv.org/abs/2412.19437) that enables models to predict multiple future tokens simultaneously during pre-training. Instead of learning to predict only the next token at each position, MTP adds auxiliary prediction heads that predict tokens 2, 3, or more positions ahead. + +### Key Benefits + +- **Densified Training Signals**: Multiple learning signals per training iteration improve data efficiency +- **Pre-Planning Representations**: Models learn internal representations that encode information about future tokens +- **Speculative Decoding Foundation**: MTP-trained models can serve as foundation for faster inference via speculative decoding + +### When to Use MTP + +MTP is most beneficial for: + +- **Large-scale pre-training** (models > 10B parameters) +- **Data-constrained scenarios** where maximizing learning from limited data is critical +- **Training foundation models** intended for downstream fine-tuning or speculative decoding + +MTP is primarily used for pre-training. + +### Additional Resources + +- [DeepSeek-V3 Technical Report](https://arxiv.org/abs/2412.19437) - Original paper introducing MTP +- [DeepSeek-V3 GitHub](https://github.com/deepseek-ai/DeepSeek-V3) - Official implementation +- [Megatron Core MTP Guide](https://github.com/NVIDIA/Megatron-LM/blob/main/docs/user-guide/features/multi_token_prediction.md) - Low-level implementation details + +## Configuration Parameters + +MTP is controlled by two primary parameters: + +| Parameter | Type | Default | Description | Typical Range | +|-----------|------|---------|-------------|---------------| +| `mtp_num_layers` | int | `None` (disabled) | Number of auxiliary prediction depths. Each layer predicts tokens N positions ahead (N=1,2,...,mtp_num_layers). | 1-2 | +| `mtp_loss_scaling_factor` | float | `0.1` | Weight applied to MTP losses relative to main next-token loss. Controls the contribution of auxiliary predictions to the total loss. | 0.05-0.2 | + +### Loss Calculation + +The total training loss combines the main next-token prediction loss with averaged MTP losses: + +``` +total_loss = main_loss + (avg_mtp_loss * mtp_loss_scaling_factor) + +where: + avg_mtp_loss = mean([mtp_1_loss, mtp_2_loss, ..., mtp_N_loss]) +``` + +### Parameter Tuning Guidelines + +**`mtp_num_layers`:** +- Start with `1` for most models (predicts 1 token ahead) +- Use `2` for models > 100B parameters if memory allows +- Higher values increase memory usage and training time proportionally + +**`mtp_loss_scaling_factor`:** +- Default `0.1` works well for most models +- Increase to `0.15-0.2` if MTP losses aren't decreasing +- Decrease to `0.05-0.08` if main loss is being overshadowed +- Scale factor should be proportional to `mtp_num_layers` (more layers → lower factor) + +## Basic Usage: Training from Scratch + +### Minimal Configuration Example + +Here's a minimal example using the Qwen3 30B-A3B recipe with MTP enabled: + +```python +from megatron.bridge.recipes.qwen.qwen3_moe import qwen3_30b_a3b_pretrain_config +from megatron.bridge.training.pretrain import pretrain +from megatron.bridge.training.gpt_step import forward_step +from megatron.bridge.training.config import ConfigContainer + +log_dir = "/path/to/log/dir" +cfg: ConfigContainer = qwen3_30b_a3b_pretrain_config() +cfg.logger.tensorboard_dir = log_dir + "/tb_logs" +cfg.checkpoint.save = log_dir + "/checkpoints" +cfg.checkpoint.load = log_dir + "/checkpoints" +# Set up training dataset +cfg.dataset.blend=[[ + f"/path/to/dclm/preprocessed/dclm_{i:02d}_text_document" + for i in range(1, 11) +], None] +cfg.dataset.split="9999,8,2" +cfg.dataset.path_to_cache = "/path/to/cache" +# cfg.model.num_layers = 8 # train a smaller model if OOM +# MTP Configuration +cfg.model.mtp_num_layers = 1 +cfg.model.mtp_loss_scaling_factor = 0.1 +pretrain(cfg, forward_step) +``` +Follow the [DCLM Tutorial](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/tutorials/data/dclm) to prepare the training data + +## MTP with Pipeline Parallelism + +When using Pipeline Parallelism (PP), **MTP layers must be placed in the last pipeline stage** alongside the loss computation layer. Configure this using custom pipeline layout settings (`pipeline_model_parallel_split_rank`). + +### Pipeline Layout Guidelines + +MTP layers take approximately the same training time as a regular transformer layer. When configuring your pipeline layout: + +- **Place MTP in the last PP stage** (required for correct loss computation) +- **Reduce layers in other PP ranks** to balance computation time across stages +- Example: For a 21-layer model with PP=4 and `mtp_num_layers=1`, you might use splits like `[5, 6, 6, 4]` instead of `[5, 5, 5, 6]` to account for MTP overhead in the last stage + +## Parallelism Support + +MTP is compatible with all major parallelism strategies in Megatron-Bridge: + +| Parallelism Type | Support Status | Notes | +|------------------|----------------|-------| +| **Tensor Parallelism (TP)** | ✅ Fully Supported | MTP layers are automatically sharded across TP ranks | +| **Pipeline Parallelism (PP)** | ✅ Supported with Constraint | MTP must be in last pipeline stage (see above) | +| **Expert Parallelism (EP)** | ✅ Fully Supported | Works with MoE models (DeepSeek-V3, Mixtral, etc.) | +| **Context Parallelism (CP)** | ✅ Fully Supported | MTP supports long-context training via CP | +| **Data Parallelism (DP)** | ✅ Fully Supported | Standard data parallelism works transparently | + +## Monitoring MTP Training + +### Per-Layer Loss Logging + +During training, you'll see losses for each MTP depth logged separately: + +``` +iteration 100/ 300000 | consumed samples: 3200 | elapsed time per iteration (ms): 3738.6 | learning rate: 6.000000E-05 | global batch size: 32 | lm loss: 7.968678E+00 | load_balancing_loss: 1.329517E+00 | mtp_1 loss: 7.925096E+00 | loss scale: 1.0 | grad norm: 1.040 | number of skipped iterations: 0 | number of nan iterations: 0 | +``` + +### Interpreting Loss Values + +![MTP Loss Curves](/assets/training/images/mtp_loss.png) + +The figure above shows typical training curves for MTP-enabled training: +- **Left**: MTP auxiliary loss (`mtp_1 loss`) tracking the first additional token prediction +- **Right**: Main language model loss (`lm loss`) for standard next-token prediction + +**Expected Patterns:** + +- **MTP losses are higher than main loss**: Predicting tokens further in the future is inherently harder. In the example above, `mtp_1 loss` (~4.3) is higher than `lm loss` (~3.9) at 3500 iterations. + +- **All losses decrease over training**: Both main and MTP losses should trend downward, as shown in the curves above. + +- **Loss gap remains relatively stable**: The difference between main and MTP losses should not grow significantly over training. + +**Red Flags:** + +- **NaN values**: Indicates training instability (see Troubleshooting section) +- **Diverging losses**: If MTP losses increase while main loss decreases, reduce `mtp_loss_scaling_factor` +- **Widening gap**: If MTP losses fall behind by > 1.0, increase `mtp_loss_scaling_factor` + +**MTP vs Non-MTP Comparison:** + +![MTP Loss Comparison](/assets/training/images/mtp_loss_comparison.png) + +The figure above compares `lm loss` between MTP-enabled (blue) and non-MTP (red) training runs on Qwen3-30B-A3B. The curves do not differ significantly in the first few thousand iterations. Notably, the MTP-enabled run shows smoother behavior around iterations 1000 and 2300, where the non-MTP run exhibits more pronounced spikes. + +### TensorBoard/WandB Visualization + +MTP losses are automatically logged to TensorBoard and/or WandB. Look for: + +- `lm loss` - Main next-token prediction loss +- `mtp_1 loss` - First auxiliary prediction loss +- `mtp_2 loss` - Second auxiliary prediction loss (if `mtp_num_layers=2`) + +### Training Characteristics + +- MTP adds computational overhead due to additional forward passes +- Memory usage increases proportionally to `mtp_num_layers` +- MTP is designed to improve data efficiency during pre-training + +**Model Performance:** + +- MTP provides additional training signals at each token position +- Can potentially improve downstream task performance +- MTP-trained models can be used for speculative decoding during inference + +## Current Limitations + +The following features are not yet supported with MTP: + +| Feature | Status | Workaround | +|---------|--------|------------| +| **HuggingFace ↔ Megatron Checkpoint Conversion** | ⚠️ Model-specific | Conversion support varies by model; check model-specific documentation | +| **Sequence Packing (Fine-Tuning)** | ❌ Not supported | For pre-training, no issues. For fine-tuning, set `packed_sequence_specs=None` | +| **Cross-Attention** | ❌ Not supported | MTP only works with decoder-only models (GPT, Llama, etc.) | +| **Learned Absolute Position Embeddings** | ❌ Not supported | Use RoPE (rotary position embeddings) or no position embeddings | +| **Block-Based Activation Recomputation** | ❌ Not supported | Use `recompute_granularity="selective"` or `"uniform"` | + +### Important Notes + +**Checkpoint Conversion:** + +HuggingFace ↔ Megatron checkpoint conversion with MTP is model-specific. Some models have conversion support planned, while others may not support MTP parameter mapping. Check the documentation for your specific model. + +**Sequence Packing:** + +MTP is incompatible with fine-tuning sequence packing (e.g., SFT with packed sequences). For pre-training, there are no sequence packing restrictions. + +## Troubleshooting Guide + +### Error: Out of Memory (OOM) + +MTP increases memory usage proportionally to `mtp_num_layers`. Try: +- Reduce `mtp_num_layers` to 1 +- Enable activation recomputation: `recompute_granularity="selective"` +- Increase pipeline parallelism +- Reduce micro batch size + +### Error: MTP Loss is NaN + +Training instability. Try: +- Lower learning rate +- Enable gradient clipping: `clip_grad=1.0` +- Use BF16 instead of FP16 +- Reduce `mtp_loss_scaling_factor` to 0.05 + +### Expected Log: `MTP layers not found on this PP rank` + +This is normal. Only the last pipeline stage builds MTP layers. + +## Additional Resources + +### Code Examples + +- **DeepSeek-V3 Recipe**: [`src/megatron/bridge/recipes/deepseek/deepseek_v3.py`](/../../src/megatron/bridge/recipes/deepseek/deepseek_v3.py) + - Example of MTP with large-scale MoE model + - Predefined pipeline layouts for PP + MTP + +- **Qwen3-Next Recipe**: [`src/megatron/bridge/recipes/qwen/qwen3_next.py`](/../../src/megatron/bridge/recipes/qwen/qwen3_next.py) + - Clean example of MTP configuration for dense models + - Good starting point for custom recipes + +- **MTP Core Implementation**: [`3rdparty/Megatron-LM/megatron/core/transformer/multi_token_prediction.py`](/../../3rdparty/Megatron-LM/megatron/core/transformer/multi_token_prediction.py) + - Low-level MTP layer implementation + - Loss computation and logging helpers + +### Documentation + +- **Megatron Core MTP Guide**: [`3rdparty/Megatron-LM/docs/user-guide/features/multi_token_prediction.md`](https://github.com/NVIDIA/Megatron-LM/blob/main/docs/user-guide/features/multi_token_prediction.md) + - Implementation notes and design decisions + +- **Pipeline Parallelism Guide**: [`docs/parallelisms.md`](/../parallelisms) + - Understanding pipeline parallelism layouts + - Best practices for PP configuration + +### External Resources + +- **DeepSeek-V3 Technical Report**: [https://arxiv.org/abs/2412.19437](https://arxiv.org/abs/2412.19437) + - Original paper introducing MTP + - Section 3.2: "Multi-Token Prediction" + - Training details and ablation studies + +- **DeepSeek-V3 GitHub**: [https://github.com/deepseek-ai/DeepSeek-V3](https://github.com/deepseek-ai/DeepSeek-V3) + - Official implementation and model weights + - Training configurations and hyperparameters + +- **Megatron-LM GitHub**: [https://github.com/NVIDIA/Megatron-LM](https://github.com/NVIDIA/Megatron-LM) + - Upstream Megatron-Core implementation + - Issues and discussions + +### Getting Help + +If you encounter issues not covered in this guide: + +1. Check the [Megatron-Bridge GitHub Issues](https://github.com/NVIDIA-NeMo/Megatron-Bridge/issues) +2. Review the [Megatron-LM Discussions](https://github.com/NVIDIA/Megatron-LM/discussions) + +When reporting issues, include: +- Full training configuration (recipe and parameters) +- Error messages and stack traces +- GPU type and count +- Megatron-Core version (`pip show megatron-core`) diff --git a/fern/v0.2.0/pages/training/optimizer-scheduler.mdx b/fern/v0.2.0/pages/training/optimizer-scheduler.mdx new file mode 100644 index 0000000000..169a7682d8 --- /dev/null +++ b/fern/v0.2.0/pages/training/optimizer-scheduler.mdx @@ -0,0 +1,119 @@ +--- +title: "Optimizer and Scheduler Configuration" +description: "" +--- + +The optimizer and scheduler configurations control optimization algorithms, learning rate schedules, and weight decay strategies. + +## OptimizerConfig (from Megatron Core) + +The `OptimizerConfig` contains all parameters for the optimization algorithm and comes directly from Megatron Core. Key parameters include: + +| Parameter | Type | Description | +|-----------|------|-------------| +| `optimizer` | `str` | Optimizer type ("adam", "sgd", etc.) | +| `lr` | `float` | Base learning rate | +| `min_lr` | `float` | Minimum learning rate for decay schedules | +| `weight_decay` | `float` | L2 regularization coefficient | +| `adam_beta1` | `float` | Adam optimizer beta1 parameter | +| `adam_beta2` | `float` | Adam optimizer beta2 parameter | +| `adam_eps` | `float` | Adam optimizer epsilon parameter | +| `clip_grad` | `float` | Gradient clipping threshold | +| `use_distributed_optimizer` | `bool` | Enable distributed optimizer for memory efficiency | +| `overlap_grad_reduce` | `bool` | Overlap gradient reduction with computation | +| `overlap_param_gather` | `bool` | Overlap parameter gathering with computation | +| `bf16` | `bool` | Use BF16 precision for training | +| `fp16` | `bool` | Use FP16 precision for training | + +## SchedulerConfig + +The `SchedulerConfig` controls learning rate scheduling and weight decay progression throughout training. + +### Learning Rate Scheduling + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `lr_decay_style` | `Literal["constant", "linear", "cosine", "inverse-square-root", "WSD"]` | `"linear"` | Learning rate decay function | +| `lr_decay_iters` | `Optional[int]` | `None` | Iterations to decay LR over (defaults to `train_iters`). Use for iteration-based training. | +| `lr_decay_samples` | `Optional[int]` | `None` | Samples to decay LR over (defaults to `train_samples`). Use for sample-based training. | +| `lr_warmup_iters` | `int` | `0` | Iterations to linearly warmup learning rate. Use for iteration-based training. | +| `lr_warmup_samples` | `int` | `0` | Samples to linearly warmup learning rate. Use for sample-based training. | +| `lr_warmup_fraction` | `Optional[float]` | `None` | Fraction of decay iterations/samples to use for warmup (works with both modes) | +| `lr_warmup_init` | `float` | `0.0` | Initial learning rate for warmup phase | + +**Scheduler Mode Selection** + +The scheduler supports two modes that must align with your training configuration: + +1. **Iteration-based scheduling**: Use `lr_decay_iters` and `lr_warmup_iters` with `train_iters`. +2. **Sample-based scheduling**: Use `lr_decay_samples` and `lr_warmup_samples` with `train_samples`. + +**Important constraints** +- Cannot mix iteration-based and sample-based scheduler parameters. +- Your scheduler mode must match your training mode (iteration-based vs sample-based). +- `lr_warmup_fraction` is compatible with both modes but cannot be used with explicit warmup iterations/samples. + +### WSD (Warmup-Stable-Decay) Scheduling + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `lr_wsd_decay_style` | `Literal["exponential", "linear", "cosine"]` | `"exponential"` | Decay style for WSD annealing phase | +| `lr_wsd_decay_iters` | `Optional[int]` | `None` | Iterations for WSD annealing phase. Use for iteration-based training. | +| `lr_wsd_decay_samples` | `Optional[int]` | `None` | Samples for WSD annealing phase. Use for sample-based training. | + +### Weight Decay Scheduling + +Parameters for controlling the progression of weight decay during training, including start and end values and the scheduling strategy: + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `start_weight_decay` | `Optional[float]` | `None` | Initial weight decay coefficient | +| `end_weight_decay` | `Optional[float]` | `None` | Final weight decay coefficient | +| `weight_decay_incr_style` | `Literal["constant", "linear", "cosine"]` | `"constant"` | Weight decay progression style | + +### Checkpoint Integration + +Parameters for managing how scheduler settings are applied during checkpoint loading, allowing control over whether to prioritize config values or restore from saved state: + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `override_opt_param_scheduler` | `bool` | `False` | Reset scheduler values from config, ignoring checkpoint | +| `use_checkpoint_opt_param_scheduler` | `bool` | `False` | Use scheduler values from checkpoint, ignoring config | + +### Computed Fields + +These fields are automatically calculated during configuration validation and help align training schedules with the configured batch size and iteration counts: + +| Field | Description | +|-------|-------------| +| `lr_warmup_steps` | Total steps for warmup (calculated from iterations and batch size) | +| `lr_decay_steps` | Total steps for decay (calculated from iterations and batch size) | +| `wd_incr_steps` | Total steps for weight decay progression | +| `wsd_decay_steps` | Total steps for WSD annealing phase | + +## Learning Rate Schedules + +The following scheduling strategies define how the learning rate evolves during training, each suited to different convergence behaviors and model types: +| Schedule Type | Description | +|-------------------------|-----------------------------------------------------------------------------| +| **Constant** | Learning rate remains fixed throughout training. | +| **Linear** | Learning rate decreases linearly from the base LR to the minimum LR. | +| **Cosine** | Learning rate follows a cosine decay curve from base LR to minimum LR. | +| **Inverse Square Root** | Learning rate decays proportionally to the inverse square root of the step. | + +## WSD (Warmup-Stable-Decay) +The WSD schedule divides learning rate progression into three distinct phases, offering fine-grained control over early ramp-up, mid-training stability, and final decay: +| Phase | Description | +|-----------|----------------------------------------------------------| +| **Warmup** | Learning rate increases linearly from initial value to base LR. | +| **Stable** | Learning rate remains constant at base LR. | +| **Decay** | Learning rate decays to minimum LR using a specified style (e.g., exponential, linear, cosine). | + +## Weight Decay Scheduling + +These scheduling options control how the weight decay coefficient changes over time, allowing for regularization strategies that adapt to different training phases: +| Schedule Type | Description | +|---------------|-----------------------------------------------------------------------------| +| **Constant** | Fixed weight decay throughout training. | +| **Linear** | Linear progression from start to end weight decay. | +| **Cosine** | Cosine progression from start to end weight decay. | diff --git a/fern/v0.2.0/pages/training/packed-sequences.mdx b/fern/v0.2.0/pages/training/packed-sequences.mdx new file mode 100644 index 0000000000..939ad62f8c --- /dev/null +++ b/fern/v0.2.0/pages/training/packed-sequences.mdx @@ -0,0 +1,187 @@ +--- +title: "Packed Sequences" +description: "" +--- + +This guide explains how to use packed sequences in Megatron Bridge for efficient supervised fine-tuning (SFT) and parameter-efficient fine-tuning (PEFT). + +## Overview + +When fine-tuning large language models, GPU under-utilization often occurs due to inefficient input data structure. This inefficiency arises because many fine-tuning datasets have a skewed distribution of sequence lengths, with many short sequences and a few long ones, following [Zipf's Law](https://en.wikipedia.org/wiki/Zipf%27s_law). Since transformer models require fixed-length inputs, shorter sequences must be padded with many padding tokens. + +This leads to two main inefficiencies: + +- Computation performed on the pad tokens is eventually masked out, resulting in wasted GPU computation. +- Micro batch size is often limited by the batch which contains longer sequences, so that most other micro batches have under-utilized GPU memory. + +Packed sequences is a training technique where multiple training sequences (examples) are concatenated into one long sequence (pack). This technique greatly reduces the number of padding tokens, allowing more meaningful tokens to be processed in each micro batch. As a result, it maximizes both GPU compute and GPU memory utilization. + +**Note:** Sequence packing is primarily beneficial for fine-tuning workloads. Megatron-style pretraining datasets (using `IndexedDataset` and `GPTDataset`) already concatenate documents during sampling to fill sequences to the target length, eliminating padding tokens without requiring the boundary-aware packing infrastructure described here. For supervised fine-tuning, however, naive concatenation is insufficient—each training example must be treated individually to preserve data quality. + +The conventional solution is to build a custom attention mask (specifically, a block triangular mask) to mask out attention values between sequences. However, this increases the complexity of attention from $\sum_i {s_i}^2$ to $\Big({\sum_i {s_i}}\Big)^2$, where $s_i$ is the length of the $i$th subsequence. In practice, the conventional solution puts a limit on the packed sequence size. + +Instead, Megatron Bridge provides a highly optimized version of sequence packing which makes use of variable-length attention kernels in FlashAttention and TransformerEngine. Instead of providing a custom attention mask, information about sequence boundaries is passed in with the `cu_seqlens` variable (short for cumulative sequence length). With this approach, attention values between sequences are never calculated, so the complexity of attention remains at $\sum_i {s_i}^2$. This allows the packed sequence size to increase to arbitrary lengths without affecting the memory complexity, so that GPU memory can be fully utilized. + +The packed sequence implementation automatically creates [`bridge.data.datasets.sft.GPTSFTPackedDataset`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.data.datasets.sft.GPTSFTPackedDataset.html) instances when `.npy` files are detected, providing optimized data loading and batching for packed sequences. + +## Using Packed Sequences + +### Prepare the Dataset + +In Megatron Bridge, the packed dataset is automatically prepared before training using the [`bridge.data.datasets.packed_sequence.prepare_packed_sequence_data`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.data.datasets.packed_sequence.prepare_packed_sequence_data.html) function, eliminating the need for any additional preprocessing steps. + +### Configure Packed Sequences + +Packed sequences are configured through the [`bridge.training.config.FinetuningDatasetConfig`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.FinetuningDatasetConfig.html) by specifying `packed_sequence_specs`: + +```python +from megatron.bridge.training.config import ConfigContainer, FinetuningDatasetConfig +from megatron.bridge.data.datasets.packed_sequence import PackedSequenceSpecs + +config = ConfigContainer( + # ... other configurations + dataset=FinetuningDatasetConfig( + dataset_root="/path/to/your/dataset", + seq_length=2048, + packed_sequence_specs=PackedSequenceSpecs( + packed_sequence_size=2048, + tokenizer_model_name="your_tokenizer_name", + ), + ), + # ... other configurations +) +``` + +### PackedSequenceSpecs Configuration + +The [`bridge.data.datasets.packed_sequence.PackedSequenceSpecs`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.data.datasets.packed_sequence.PackedSequenceSpecs.html) class provides the following configuration options: + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `packed_sequence_size` | `int` | `-1` | If positive, enables sequence packing with the specified pack size. If ≤ 0, sequence packing is disabled. | +| `tokenizer_model_name` | `str` | `None` | Tokenizer model name for tracking, since different tokenizers produce different packed datasets. | +| `packed_train_data_path` | `str` | `None` | Custom path for packed training dataset file (`.npy` format). | +| `packed_val_data_path` | `str` | `None` | Custom path for packed validation dataset file (`.npy` format). | +| `packed_metadata_path` | `str` | `None` | Custom path for packing metadata file (`.jsonl` format). | +| `pad_seq_to_mult` | `int \| None` | `None` | Pad each sample to a multiple of this value when generating packed datasets (e.g., set to `2 * context_parallel_size` for THD CP). | +| `pad_cu_seqlens` | `bool` | `False` | Whether to pad `cu_seqlens` to constant size, required for CUDA graphs. | + +### Batch Size Considerations + +When using packed sequences, you must adjust your batch sizes: + +1. **Micro batch size must be set to 1**: This constraint arises because samples in a micro batch are no longer stacked; they are now concatenated during the data preparation step. Consequently, micro batch size becomes irrelevant when using packed sequences. + +2. **Global batch size must be adjusted**: Since each pack now contains multiple sequences, the global batch size needs to be reduced by the average number of sequences per pack `n` where `n = num_sequences_in_dataset / num_packs` (equivalently, `n = packed_sequence_size / average_seq_len`). This ensures that each gradient iteration sees, on average, the same number of tokens. The value of `n` is printed out during the data preparation step. You may need to run training once, obtain the value of `n` from the logs, then run your training script again with the updated global batch size. + +### Full Configuration Example + +```python +from megatron.bridge.training.config import ( + ConfigContainer, TrainingConfig, CheckpointConfig, SchedulerConfig +) +from megatron.bridge.training.config import FinetuningDatasetConfig +from megatron.bridge.data.datasets.packed_sequence import PackedSequenceSpecs +from megatron.bridge.peft.lora import LoRA +from megatron.core.optimizer import OptimizerConfig + +config = ConfigContainer( + model=model_provider, + train=TrainingConfig( + train_iters=1000, + global_batch_size=32, # Reduced from original due to packing + micro_batch_size=1, # Required for packed sequences + eval_interval=100, + ), + optimizer=OptimizerConfig( + optimizer="adam", + lr=1e-4, + weight_decay=0.01, + bf16=True, + use_distributed_optimizer=True, + ), + scheduler=SchedulerConfig( + lr_decay_style="cosine", + lr_warmup_iters=100, + lr_decay_iters=1000, + ), + dataset=FinetuningDatasetConfig( + dataset_root="/path/to/dataset", + seq_length=2048, + packed_sequence_specs=PackedSequenceSpecs( + packed_sequence_size=2048, + tokenizer_model_name="llama2_tokenizer", + ), + ), + checkpoint=CheckpointConfig( + pretrained_checkpoint="/path/to/pretrained/model", + save="/path/to/checkpoints", + save_interval=200, + ), + peft=LoRA( + target_modules=["linear_qkv", "linear_proj", "linear_fc1", "linear_fc2"], + dim=16, + alpha=32, + dropout=0.1, + ), + # ... other configurations +) +``` + +## File Organization + +When using packed sequences, the [`bridge.data.builders.finetuning_dataset.FinetuningDatasetBuilder`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.data.builders.finetuning_dataset.FinetuningDatasetBuilder.html) automatically organizes files in your dataset directory: + +``` +dataset_root/ +├── training.jsonl # Original training data +├── validation.jsonl # Original validation data +└── packed/ + └── {tokenizer_name}/ + ├── training_{packed_size}.npy # Packed training data + ├── validation_{packed_size}.npy # Packed validation data + └── {packed_size}_metadata.jsonl # Packing metadata +``` + +The tokenizer name and packed sequence size are automatically incorporated into the file paths to avoid conflicts when using different configurations. + +## Advanced Configuration + +### Custom File Paths + +You can specify custom paths for packed data files: + +```python +packed_sequence_specs = PackedSequenceSpecs( + packed_sequence_size=4096, + tokenizer_model_name="custom_tokenizer", + packed_train_data_path="/custom/path/training_packed.npy", + packed_val_data_path="/custom/path/validation_packed.npy", + packed_metadata_path="/custom/path/metadata.jsonl", +) +``` + +### CUDA Graphs Support + +For CUDA graphs compatibility, enable `pad_cu_seqlens`: + +```python +packed_sequence_specs = PackedSequenceSpecs( + packed_sequence_size=2048, + pad_cu_seqlens=True, # Required for CUDA graphs + tokenizer_model_name="your_tokenizer", +) +``` + +When `pad_cu_seqlens=True`, you must also set `pad_to_max_length=True` in your dataset configuration. + +## API Reference + +For detailed API documentation, see: + +- [`bridge.training.config.FinetuningDatasetConfig`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.FinetuningDatasetConfig.html) - Main dataset configuration class +- [`bridge.data.datasets.packed_sequence.PackedSequenceSpecs`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.data.datasets.packed_sequence.PackedSequenceSpecs.html) - Packed sequence configuration +- [`bridge.data.datasets.packed_sequence.prepare_packed_sequence_data`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.data.datasets.packed_sequence.prepare_packed_sequence_data.html) - Data preparation function +- [`bridge.data.datasets.sft.GPTSFTPackedDataset`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.data.datasets.sft.GPTSFTPackedDataset.html) - Packed sequence dataset implementation +- [`bridge.data.builders.finetuning_dataset.FinetuningDatasetBuilder`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.data.builders.finetuning_dataset.FinetuningDatasetBuilder.html) - Dataset builder with packing support +- [`bridge.training.gpt_step.get_packed_seq_params`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.gpt_step.get_packed_seq_params.html) - Packed sequence parameter extraction for training diff --git a/fern/v0.2.0/pages/training/peft.mdx b/fern/v0.2.0/pages/training/peft.mdx new file mode 100644 index 0000000000..ef14d3e602 --- /dev/null +++ b/fern/v0.2.0/pages/training/peft.mdx @@ -0,0 +1,312 @@ +--- +title: "Parameter-Efficient Fine-Tuning (PEFT)" +description: "" +--- + +This guide explains how to configure and use PEFT in Megatron Bridge—covering LoRA and DoRA, required checkpoints, example configurations, and the internal design and training workflow—so you can integrate, scale, and checkpoint adapters efficiently. + +## Model Customization +Customizing models enables you to adapt a general pre-trained model to a specific use case or domain. This process produces a fine-tuned model that retains the broad knowledge from pretraining while delivering more accurate outputs for targeted downstream tasks. + +Model customization is typically achieved through supervised fine-tuning, which falls into two main approaches: Full-Parameter Fine-Tuning, known as Supervised Fine-Tuning (SFT), and Parameter-Efficient Fine-Tuning (PEFT). + +In SFT, all model parameters are updated to align the model’s outputs with the task-specific requirements. This approach often yields the highest performance but can be computationally intensive. + +PEFT, by contrast, updates only a small subset of parameters that are inserted into the base model at strategic locations. The base model weights remain frozen, and only the adapter modules are trained. This significantly reduces the number of trainable parameters—often to less than 1%—while still achieving near-SFT levels of accuracy. + +As language models continue to grow in size, PEFT is gaining popularity for its efficiency and minimal hardware demands, making it a practical choice for many real-world applications. + +## PEFT Configuration + +PEFT is configured as an optional attribute in `ConfigContainer`: + +```python +from megatron.bridge.training.config import ConfigContainer +from megatron.bridge.peft.lora import LoRA + +config = ConfigContainer( + # ... other required configurations + peft=LoRA( + target_modules=["linear_qkv", "linear_proj", "linear_fc1", "linear_fc2"], + dim=16, + alpha=32, + dropout=0.1, + ), + checkpoint=CheckpointConfig( + pretrained_checkpoint="/path/to/pretrained/checkpoint", # Required for PEFT + save="/path/to/peft/checkpoints", + ), +) +``` + + +**Requirements**: PEFT requires `checkpoint.pretrained_checkpoint` to be set to load the base model weights. + + +## Supported PEFT Methods + +### [LoRA: Low-Rank Adaptation of Large Language Models](https://arxiv.org/abs/2106.09685) + +LoRA makes fine-tuning efficient by representing weight updates with two low-rank decomposition matrices. The original model weights remain frozen, while the low-rank decomposition matrices are updated to adapt to the new data, keeping the number of trainable parameters low. In contrast with adapters, the original model weights and adapted weights can be combined during inference, avoiding any architectural change or additional latency in the model at inference time. + +In Megatron Bridge, you can configure both the adapter bottleneck dimension and the target modules where LoRA is applied. LoRA supports any linear layer, which in transformer models typically includes: + +1. Query, key, and value (QKV) attention projections +2. The attention output projection +3. One or both MLP layers + +Megatron Bridge fuses the QKV projections into a single linear layer. As a result, LoRA learns a unified low-rank adaptation for the combined QKV representation. + +```python +from megatron.bridge.peft.lora import LoRA + +lora_config = LoRA( + target_modules=["linear_qkv", "linear_proj", "linear_fc1", "linear_fc2"], + dim=16, # Rank of adaptation + alpha=32, # Scaling parameter + dropout=0.1, # Dropout rate +) +``` + +#### Key Parameters +The following table lists key hyperparameters for configuring DoRA, which control its module targeting, adaptation rank, scaling behavior, and regularization strategy. +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `target_modules` | `List[str]` | All linear layers | Modules to apply DoRA to | +| `dim` | `int` | `32` | Rank of the low-rank adaptation | +| `alpha` | `float` | `16` | Scaling parameter for DoRA | +| `dropout` | `float` | `0.0` | Dropout rate for DoRA layers | + +#### Target Modules +The following table lists specific submodules within transformer architectures that are commonly targeted for LoRA, enabling efficient fine-tuning of attention and feedforward components: +| Module | Description | +|---------------|---------------------------------------------| +| `linear_qkv` | Query, key, value projections in attention | +| `linear_proj` | Attention output projection | +| `linear_fc1` | First MLP layer | +| `linear_fc2` | Second MLP layer | + +#### Wildcard Target Modules +For more granular targeting, individual layers can be targeted for the adapters. +```python +# Target specific layers only +lora_config = LoRA( + target_modules=[ + "*.layers.0.*.linear_qkv", # First layer only + "*.layers.1.*.linear_qkv", # Second layer only + ] +) +``` + +### Canonical LoRA: Performant vs Canonical Variants + +There are two variants of LoRA implemented in Megatron Bridge: "performant LoRA" (`LoRA`) and "canonical LoRA" (`CanonicalLoRA`). + +The distinction comes from the fact that Megatron Core optimizes the implementation of the following two linear modules by fusing multiple linear layers into one layer. When these layers are adapted with LoRA, the performant version also uses only one adapter for the linear module. The two linear modules are: + +1. `linear_qkv`: The projection matrix in self attention that transforms hidden state to query, key and value. Megatron Core fuses these three projection matrices into a single matrix to efficiently parallelize the matrix multiplication. Hence, performant LoRA applies a single adapter to the qkv projection matrix, whereas canonical LoRA applies three adapters. +2. `linear_fc1`: The first linear layer in the MLP module before the intermediate activation. For gated linear activations, Megatron Core fuses the up and gate projection matrices into a single matrix for efficient parallelization. Hence, performant LoRA applies a single adapter to the up and gate projection matrices, whereas canonical LoRA applies two adapters. + +The following two figures illustrate the difference between canonical and performant LoRA, using the `linear_qkv` layer as an example. Canonical LoRA runs three adapters sequentially, while performant LoRA runs one adapter. + +![canonical_lora.png](/assets/training/images/canonical_lora.png) + +![performant_lora.png](/assets/training/images/performant_lora.png) + +Canonical LoRA conforms more closely to reference implementations, though it is slower in comparison since it performs several matrix multiplications sequentially, as described above. Performant LoRA has fewer parameters than canonical LoRA and can often achieve the same level of accuracy as canonical LoRA. + +Though not immediately apparent, performant LoRA is mathematically equivalent to canonical LoRA when the $A_q$, $A_k$, $A_v$ matrices are tied (i.e. forced to share the same weight during training) in `linear_qkv`, and similarly when the $A_{up}$, $A_{gate}$ matrices are tied in `linear_fc1`. + + +Let $[x \quad y]$ denote matrix concatenation. (In Megatron Bridge, this concatenation is done in an interleaved fashion, but this does not affect the proof below.) + +Let $A_q = A_k = A_v = A_{qkv}$ (weight tying) + +Then + +$$ +\begin{align} +& [query \quad key \quad value] \\ += & [W_q x + B_q A_q x \quad W_k x + B_k A_k x \quad W_v x + B_v A_v x] \quad\quad \text{(canonical formulation)} \\ += & [W_q x + B_q (A_{qkv} x) \quad W_k x + B_k (A_{qkv} x) \quad W_v x + B_v (A_{qkv} x)] \\ += & [W_q \quad W_k \quad W_v] x + [B_q \quad B_k \quad B_v]A_{qkv} x \\ += & W_{qkv} x + B_{qkv} A_{qkv} x \quad\quad \text{(performant formulation)} +\end{align} +$$ + +Note: dimensions of weight matrices are as follows: + +$$ +\begin{align} +W_q: &\ h \times n_q d \qquad & A_q: &\ h \times r \qquad & B_q: &\ r \times n_q d \\ +W_k: &\ h \times n_{kv} d \qquad & A_k: &\ h \times r \qquad & B_k: &\ r \times n_{kv} d \\ +W_v: &\ h \times n_{kv} d \qquad & A_v: &\ h \times r \qquad & B_v: &\ r \times n_{kv} d \\ +W_{qkv}: &\ h \times (n_q+2n_{kv})d \qquad & A_{qkv}: &\ h \times r \qquad & B_{qkv}: &\ r \times (n_q+2n_{kv})d +\end{align} +$$ + +Where: +- $n_q$: Number of attention heads (`num_attention_heads`). +- $n_{kv}$: Number of key value heads (`num_query_groups`). Note that if grouped query attention (GQA) is not used, $n_{kv} = n_q$. +- $h$: Transformer hidden size (`hidden_size`). +- $d$: Transformer head dimension (`kv_channels`). +- $r$: LoRA rank. + + +#### Using Canonical LoRA + +```python +from megatron.bridge.peft.canonical_lora import CanonicalLoRA + +canonical_lora_config = CanonicalLoRA( + target_modules=[ + "linear_q", "linear_k", "linear_v", # Individual Q, K, V projections + "linear_proj", # Attention output projection + "linear_fc1_up", "linear_fc1_gate", # Individual up and gate projections + "linear_fc2" # Second MLP layer + ], + dim=16, # Rank of adaptation + alpha=32, # Scaling parameter + dropout=0.1, # Dropout rate +) +``` + +#### Key Parameters + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `target_modules` | `List[str]` | All canonical linear layers | Modules to apply canonical LoRA to | +| `dim` | `int` | `32` | Rank of the low-rank adaptation | +| `alpha` | `float` | `32` | Scaling parameter for LoRA | +| `dropout` | `float` | `0.0` | Dropout rate for LoRA layers | +| `dropout_position` | `Literal["pre", "post"]` | `"pre"` | Position for applying dropout | +| `lora_A_init_method` | `str` | `"xavier"` | Initialization method for LoRA A matrix | +| `lora_B_init_method` | `str` | `"zero"` | Initialization method for LoRA B matrix | + +#### Target Modules for Canonical LoRA + +The following table lists specific submodules within transformer architectures that are targeted for canonical LoRA: + +| Module | Description | +|--------|-------------| +| `linear_q` | Query projection in attention | +| `linear_k` | Key projection in attention | +| `linear_v` | Value projection in attention | +| `linear_proj` | Attention output projection | +| `linear_fc1_up` | Up projection in MLP | +| `linear_fc1_gate` | Gate projection in MLP | +| `linear_fc2` | Second MLP layer | + + +Canonical LoRA does not support `linear_qkv` or `linear_fc1` targets. Use the individual component targets (`linear_q`, `linear_k`, `linear_v` for QKV and `linear_fc1_up`, `linear_fc1_gate` for FC1) instead. + + +### [DoRA: Weight-Decomposed Low-Rank Adaptation](https://arxiv.org/abs/2402.09353) + +DoRA decomposes the pre-trained weight into magnitude and direction. It learns a separate magnitude parameter while employing LoRA for directional updates, efficiently minimizing the number of trainable parameters. DoRA enhances both the learning capacity and training stability of LoRA, while avoiding any additional inference overhead. DoRA has been shown to consistently outperform LoRA on various downstream tasks. + +In Megatron Bridge, DoRA leverages the same adapter structure as LoRA. Megatron Bridge adds support for Tensor Parallelism and Pipeline Parallelism for DoRA, enabling DoRA to be scaled to larger model variants. + +```python +from megatron.bridge.peft.dora import DoRA + +dora_config = DoRA( + target_modules=["linear_qkv", "linear_proj", "linear_fc1", "linear_fc2"], + dim=16, # Rank of adaptation + alpha=32, # Scaling parameter + dropout=0.1, # Dropout rate +) +``` + +#### Key Parameters + +The following parameters define how LoRA is applied to your model. They control which modules are targeted, the adaptation rank, scaling behavior, and dropout configuration: + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `target_modules` | `List[str]` | All linear layers | Modules to apply DoRA to | +| `dim` | `int` | `32` | Rank of the low-rank adaptation | +| `alpha` | `float` | `16` | Scaling parameter for DoRA | +| `dropout` | `float` | `0.0` | Dropout rate for DoRA layers | + +## Full Configuration Example + +```python +from megatron.bridge.training.config import ( + ConfigContainer, TrainingConfig, CheckpointConfig +) +from megatron.bridge.data.builders.hf_dataset import HFDatasetConfig +from megatron.bridge.data.hf_processors.squad import process_squad_example +from megatron.bridge.peft.lora import LoRA +from megatron.core.optimizer import OptimizerConfig + +# Configure PEFT fine-tuning +config = ConfigContainer( + model=model_provider, + train=TrainingConfig( + train_iters=1000, + global_batch_size=64, + micro_batch_size=1, # Required for packed sequences if used + eval_interval=100, + ), + optimizer=OptimizerConfig( + optimizer="adam", + lr=1e-4, # Lower learning rate for fine-tuning + weight_decay=0.01, + bf16=True, + use_distributed_optimizer=True, + ), + scheduler=SchedulerConfig( + lr_decay_style="cosine", + lr_warmup_iters=100, + lr_decay_iters=1000, + ), + dataset=HFDatasetConfig( + dataset_name="squad", + process_example_fn=process_squad_example, + seq_length=512, + ), + checkpoint=CheckpointConfig( + pretrained_checkpoint="/path/to/pretrained/model", # Required + save="/path/to/peft/checkpoints", + save_interval=200, + ), + peft=LoRA( + target_modules=["linear_qkv", "linear_proj", "linear_fc1", "linear_fc2"], + dim=16, + alpha=32, + dropout=0.1, + ), + # ... other configurations +) +``` + +## PEFT Design in Megatron Bridge + +This section describes the internal design and architecture for how PEFT is integrated into Megatron Bridge. + +### Architecture Overview + +The PEFT framework introduces a modular design for integrating adapters into large-scale models. Its architecture consists of the following components: + +1. **Base PEFT Class**: All PEFT methods inherit from the abstract [`bridge.peft.base.PEFT`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.peft.base.PEFT.html) base class, which defines the core interface for module transformation. +2. **Module Transformation**: PEFT traverses the model structure to identify and transform target modules individually. +3. **Adapter Integration**: Adapters are injected into selected modules using a pre-wrap hook during model initialization. +4. **Checkpoint Integration**: Only adapter parameters are saved and loaded during checkpointing; base model weights remain frozen and unchanged. + +### PEFT Workflow in Training + +The training workflow for PEFT follows a structured sequence that ensures efficient fine-tuning with minimal overhead: +1. **Model Loading**: The base model is initialized from a specified pretrained checkpoint. +2. **PEFT Application**: Adapter transformations are applied after Megatron Core model initialization, but before distributed wrapping. +3. **Parameter Freezing**: Base model parameters are frozen to reduce training complexity; only adapter parameters are updated. +4. **Adapter Weight Loading**: When resuming training, adapter weights are restored from the checkpoint. +5. **Checkpoint Saving**: Only adapter states are saved, resulting in significantly smaller checkpoint files. + +### Key Benefits + +PEFT offers several advantages for scalable and efficient model fine-tuning: + +- **Reduced Checkpoint Size**: Adapter-only checkpoints are dramatically smaller than full model checkpoints. +- **Memory Efficiency**: Since gradients are computed only for adapter parameters, memory usage is significantly reduced. +- **Resume Support**: Training can be resumed seamlessly using adapter-only checkpoints, without reloading full model weights. diff --git a/fern/v0.2.0/pages/training/profiling.mdx b/fern/v0.2.0/pages/training/profiling.mdx new file mode 100644 index 0000000000..d2aeb429d8 --- /dev/null +++ b/fern/v0.2.0/pages/training/profiling.mdx @@ -0,0 +1,173 @@ +--- +title: "Profiling" +description: "" +--- + +Megatron Bridge provides built-in support for profiling training jobs using a range of performance analysis tools. These include NVIDIA Nsight Systems (Nsys) for workflow optimization, as well as PyTorch-based profilers and memory trackers to monitor performance and memory usage patterns during training. + +## ProfilingConfig Overview + +[`bridge.training.config.ProfilingConfig`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.ProfilingConfig.html) is a dataclass that encapsulates profiling-related settings for training. It resides inside the overall [`bridge.training.config.ConfigContainer`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.ConfigContainer.html), which represents the complete configuration for a training run. + +### Profiling Options + +The configuration supports two mutually exclusive profiling options: + +- **NSys profiling** (`use_nsys_profiler`) +- **PyTorch profiling** (`use_pytorch_profiler`) + +You can enable one or the other, but not both at the same time. + +### Step Range and Target Ranks + +All profiling modes allow you to configure: + +- **Step range**: `profile_step_start` and `profile_step_end` +- **Target ranks**: `profile_ranks` + +By default, profiling targets rank 0. You can specify multiple ranks to analyze different parts of your distributed training setup. + +### Advanced Profiling Features + +The configuration includes options for recording tensor shapes (`record_shapes`) and enabling memory profiling (`record_memory_history`) with a customizable output path (`memory_snapshot_path`). These features offer deeper visibility into your model’s memory consumption and tensor-level operations during training. + +## NSys Profiling + +NVIDIA Nsys is a system-wide performance analysis tool designed to help you tune and optimize CUDA applications. Megatron Bridge integrates with Nsys to enable profiling specific steps of your training job, making it easy to collect detailed performance data without manual instrumentation. + + +NSys profiling cannot be used with the `FaultTolerancePlugin` due to implementation conflicts. If both are enabled, the framework will automatically disable NSys profiling and emit a warning. + + +### Configure NSys Profiling + +Enable NSys profiling by setting `use_nsys_profiler=True` in your `ProfilingConfig`. The key configuration options include: + +```python +from megatron.bridge.training.config import ProfilingConfig + +# In your ConfigContainer setup, cfg is a ConfigContainer instance +cfg.profiling = ProfilingConfig( + use_nsys_profiler=True, + profile_step_start=10, + profile_step_end=15, + profile_ranks=[0, 1], # Profile first two ranks + record_shapes=False, # Optional: record tensor shapes +) +``` + +### Launch with NSys + +When using NSys profiling, launch your training script with the NSys command wrapper: + +```bash +nsys profile -s none -o -t cuda,nvtx --force-overwrite true --capture-range=cudaProfilerApi --capture-range-end=stop python +``` + +Replace `` with your desired output path and `` with your training script. The `--capture-range=cudaProfilerApi` option ensures profiling is controlled by the framework's step range configuration. + +### Configure Profiling with the NeMo Run NSys Plugin + +Recipe users can leverage the [`bridge.recipes.run_plugins.NsysPlugin`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.recipes.run_plugins.NsysPlugin.html) to configure NSys profiling through NeMo Run executors. The plugin provides a convenient interface for setting up profiling without manually configuring the underlying NSys command. + +```python +import nemo_run as run +from megatron.bridge.recipes.run_plugins import NsysPlugin + +# Create your recipe and executor +recipe = your_recipe_function() +executor = run.SlurmExecutor(...) + +# Configure NSys profiling via plugin +plugins = [ + NsysPlugin( + profile_step_start=10, + profile_step_end=15, + profile_ranks=[0, 1], + nsys_trace=["nvtx", "cuda"], # Optional: specify trace events + record_shapes=False, + nsys_gpu_metrics=False, + ) +] + +# Run with profiling enabled +with run.Experiment("nsys_profiling_experiment") as exp: + exp.add(recipe, executor=executor, plugins=plugins) + exp.run() +``` + +The plugin automatically configures the NSys command line options and sets up the profiling configuration in your training job. + +### Analyze Results + +After your profiling run completes, the NSys profile files (`.nsys-rep`) will be generated. To analyze them, install [NVIDIA Nsight Systems](https://developer.nvidia.com/nsight-systems) from the NVIDIA Developer website, open the files in the NSys GUI, and use the timeline view to explore the performance characteristics of your training job. + +## PyTorch Profiler + +Megatron Bridge supports the built-in PyTorch profiler, which is useful for viewing profiles in TensorBoard and understanding PyTorch-level performance characteristics. + +### Configure PyTorch Profiler + +Enable PyTorch profiling by setting `use_pytorch_profiler=True` in your `ProfilingConfig`: + +```python +from megatron.bridge.training.config import ProfilingConfig + +cfg.profiling = ProfilingConfig( + use_pytorch_profiler=True, + profile_step_start=10, + profile_step_end=15, + profile_ranks=[0], + record_shapes=True, # Record tensor shapes for detailed analysis +) +``` + +### Configure Profiling with the PyTorch Profiler Plugin + +Similar to NSys, recipe users can use the [`bridge.recipes.run_plugins.PyTorchProfilerPlugin`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.recipes.run_plugins.PyTorchProfilerPlugin.html) for convenient configuration: + +```python +from megatron.bridge.recipes.run_plugins import PyTorchProfilerPlugin + +plugins = [ + PyTorchProfilerPlugin( + profile_step_start=10, + profile_step_end=15, + profile_ranks=[0], + record_memory_history=True, + memory_snapshot_path="memory_snapshot.pickle", + record_shapes=True, + ) +] +``` + +## Memory Profiling + +Megatron Bridge provides built-in support for CUDA memory profiling to track and analyze memory usage patterns during training, including GPU memory allocation and consumption tracking. + +More information about the generated memory profiles can be found [here](https://pytorch.org/blog/understanding-gpu-memory-1/). + +### Configure Memory Profiling + +Enable memory profiling by setting `record_memory_history=True` in your `ProfilingConfig`. This can be used with either profiling mode: + +```python +from megatron.bridge.training.config import ProfilingConfig + +cfg.profiling = ProfilingConfig( + use_pytorch_profiler=True, # or use_nsys_profiler=True + profile_step_start=10, + profile_step_end=15, + profile_ranks=[0], + record_memory_history=True, + memory_snapshot_path="memory_trace.pickle", # Customize output path +) +``` + +### Analyze Memory Usage + +After the run completes, memory snapshots for each specified rank are saved to the designated path. Load these traces using the PyTorch Memory Viz tool to plot memory usage over time and detect bottlenecks or leaks in your training pipeline. + +## Optimize Profiling Accuracy + +Profiling adds overhead to your training job, so measured timings may be slightly higher than normal operation. For accurate profiling results, disable other intensive operations like frequent checkpointing during the profiled step range. Choose your profiling step range carefully to capture representative training behavior while minimizing the performance impact on the overall job. diff --git a/fern/v0.2.0/pages/training/pruning.mdx b/fern/v0.2.0/pages/training/pruning.mdx new file mode 100644 index 0000000000..b7acd86c59 --- /dev/null +++ b/fern/v0.2.0/pages/training/pruning.mdx @@ -0,0 +1,81 @@ +--- +title: "Pruning" +description: "" +--- + +Pruning reduces model size by removing redundant parameters (e.g., shrinking hidden dimensions or layers) while preserving accuracy. In Megatron Bridge, pruning is provided by [NVIDIA Model Optimizer (ModelOpt)](https://github.com/NVIDIA/Model-Optimizer) using the Minitron algorithm for GPT and Mamba-based models loaded from HuggingFace. + +## Pre-requisites + +Running the pruning example requires Megatron-Bridge and Model-Optimizer dependencies. We recommend using the NeMo container (e.g., `nvcr.io/nvidia/nemo:26.02`). To use the latest ModelOpt scripts, mount your Model-Optimizer repo to the container. + +```bash +export MODELOPT_DIR=${PWD}/Model-Optimizer # or set to your local Model-Optimizer repository path if you have cloned it +if [ ! -d "${MODELOPT_DIR}" ]; then + git clone https://github.com/NVIDIA/Model-Optimizer.git ${MODELOPT_DIR} +fi + +export DOCKER_IMAGE=nvcr.io/nvidia/nemo:26.02 +docker run \ + --gpus all \ + --shm-size=20g \ + --net=host \ + --ulimit memlock=-1 \ + --rm -it \ + -v ${MODELOPT_DIR}:/opt/Model-Optimizer \ + -v ${MODELOPT_DIR}/modelopt:/opt/venv/lib/python3.12/site-packages/modelopt \ + -w /opt/Model-Optimizer/examples/megatron_bridge \ + ${DOCKER_IMAGE} bash +``` + +Once inside the container, you need to login with your HuggingFace token to download gated datasets / models. +Note that the default dataset for pruning is [`nemotron-post-training-dataset-v2`](https://huggingface.co/datasets/nvidia/Nemotron-Post-Training-Dataset-v2), which is gated. + +```bash +huggingface-cli login --token +``` + +## Usage + +### Prune to a target parameter count (using Neural Architecture Search) + +Example: prune Qwen3-8B to 6B on 2 GPUs (Pipeline Parallelism = 2), skipping pruning of `num_attention_heads`. Defaults: 1024 samples from [nemotron-post-training-dataset-v2](https://huggingface.co/datasets/nvidia/Nemotron-Post-Training-Dataset-v2) for calibration, at most 20% depth (`num_layers`) and 40% width per prunable hyperparameter (`hidden_size`, `ffn_hidden_size`, ...), top-10 candidates evaluated for MMLU (5% sampled data) to select the best model. + +```bash +torchrun --nproc_per_node 2 prune_minitron.py \ + --pp_size 2 \ + --hf_model_name_or_path Qwen/Qwen3-8B \ + --prune_target_params 6e9 \ + --hparams_to_skip num_attention_heads \ + --output_hf_path /tmp/Qwen3-8B-Pruned-6B +``` + +### Prune to a specific architecture (using manual configuration) + +Example: prune Qwen3-8B to a fixed architecture. Defaults: 1024 samples from [nemotron-post-training-dataset-v2](https://huggingface.co/datasets/nvidia/Nemotron-Post-Training-Dataset-v2) for calibration. + +```bash +torchrun --nproc_per_node 2 prune_minitron.py \ + --pp_size 2 \ + --hf_model_name_or_path Qwen/Qwen3-8B \ + --prune_export_config '{"hidden_size": 3584, "ffn_hidden_size": 9216}' \ + --output_hf_path /tmp/Qwen3-8B-Pruned-6B-manual +``` + +To see the full list of options for advanced configurations, run: + +```bash +torchrun --nproc_per_node 1 prune_minitron.py --help +``` + +### Uneven pipeline parallelism + +If the number of layers is not divisible by the number of GPUs (pipeline parallel size), set `--num_layers_in_first_pipeline_stage` and `--num_layers_in_last_pipeline_stage`. For example, Qwen3-8B with 36 layers on 8 GPUs: set both to 3 to get 3-5-5-5-5-5-5-3 layers per GPU. + +## More information + +For more details, see the [ModelOpt pruning README](https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/megatron_bridge#readme). + +## Next steps: Knowledge Distillation + +Knowledge Distillation is required to recover the performance of the pruned model. See the [Knowledge Distillation](/distillation) guide for more details. diff --git a/fern/v0.2.0/pages/training/resiliency.mdx b/fern/v0.2.0/pages/training/resiliency.mdx new file mode 100644 index 0000000000..5fa4f2b456 --- /dev/null +++ b/fern/v0.2.0/pages/training/resiliency.mdx @@ -0,0 +1,782 @@ +--- +title: "Resiliency" +description: "" +--- + +Megatron Bridge incorporates resilient training features from the [NVIDIA Resiliency Extension](https://github.com/NVIDIA/nvidia-resiliency-ext). This extension provides fault-tolerant capabilities that help minimize downtime due to failures and interruptions during training. + +## Fault Tolerance: In Job Restart + +The fault tolerance feature can detect hangs during training and automatically restart a workload due to a hang or error. This is particularly useful when training on unreliable hardware, at very large scale, or when transient faults are common. + +### Key Features + +- **Hang Detection**: Monitors training progress and detects when ranks become unresponsive. +- **Automatic Restart**: Automatically restarts training from the last checkpoint when faults are detected. +- **Section-based Monitoring**: Uses different timeout thresholds for setup, training steps, and checkpointing operations. +- **Timeout Calculation**: Can automatically calculate optimal timeouts based on observed training behavior. +- **Multi-level Restart Logic**: Supports both in-job restarts and new job launches on failure. + +### Prerequisites + +> **Warning**: This feature is currently only supported on Slurm-based clusters. + +Before using fault tolerance features, ensure the following: + +1. **Slurm Environment**: The system must be running on a Slurm-based cluster. +2. **Checkpoint Configuration**: A valid directory for saving checkpoints must be properly configured. + +### Usage Options + +Megatron Bridge provides two ways to enable fault tolerance: + +#### Option 1: NeMo Run Plugin + +If you're using NeMo Run, the [`bridge.recipes.run_plugins.FaultTolerancePlugin`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.recipes.run_plugins.FaultTolerancePlugin.html) provides the simplest integration: + +```python +from megatron.bridge.recipes.run_plugins import FaultTolerancePlugin +import nemo_run as run + +# Configure your task +task = run.Script(...) + +# Add fault tolerance plugin +run_plugins = [ + FaultTolerancePlugin( + enable_ft_package=True, + calc_ft_timeouts=True, + num_in_job_restarts=3, + num_job_retries_on_failure=2, + initial_rank_heartbeat_timeout=1800, + rank_heartbeat_timeout=300, + ) +] + +# Run with fault tolerance +run.run(task, plugins=run_plugins, executor=executor) +``` + +#### Option 2: Direct Configuration + +If you’re a user who wants more direct control, you can configure fault tolerance manually: + +```python +from megatron.bridge.training.config import FaultToleranceConfig + +# Configure fault tolerance in your config +config.ft = FaultToleranceConfig( + enable_ft_package=True, + calc_ft_timeouts=True, + # Optional: simulate faults for testing + simulate_fault=False, + simulated_fault_type="random", +) +``` + +When directly using the configuration, you must launch your training script using the `ft_launcher` tool: + +```bash +ft_launcher \ + --rdzv_backend=c10d --rdzv_endpoint=${MASTER_ADDR}:${MASTER_PORT} \ + --nnodes=${NUM_NODES} --nproc-per-node=${NUM_GPUS_PER_NODE} \ + --ft-rank_section_timeouts=setup:600,step:180,checkpointing:420 \ + --ft-rank_out_of_section_timeout=300 \ + your_training_script.py +``` + +> **Note**: For local testing or non-Slurm environments, you must set the `GROUP_RANK` environment variable before launching `ft_launcher`: +> ```bash +> export GROUP_RANK=0 # For single-node runs +> ft_launcher ... +> ``` +> This is required because `ft_launcher` uses `use_infra_group_rank=True` by default, which expects either `SLURM_PROCID` or `GROUP_RANK` to be set. + +### Configuration Options + +The fault tolerance system can be configured through [`bridge.training.config.FaultToleranceConfig`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.FaultToleranceConfig.html): + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `enable_ft_package` | `bool` | `False` | Enable the fault tolerance package | +| `calc_ft_timeouts` | `bool` | `False` | Automatically compute optimal timeouts | +| `simulate_fault` | `bool` | `False` | Enable fault simulation for testing | +| `simulated_fault_type` | `str` | `"random"` | Type of fault to simulate: `"rank_hung"`, `"rank_killed"`, or `"random"` | +| `simulated_fault_rank` | `int` | `None` | Specific rank to simulate fault on (random if not specified) | +| `simulated_fault_base_delay` | `int` | `0` | Base delay before simulating fault | + +### Plugin Configuration Options + +When using the [`bridge.recipes.run_plugins.FaultTolerancePlugin`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.recipes.run_plugins.FaultTolerancePlugin.html), additional options are available: + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `num_in_job_restarts` | `int` | `3` | Maximum number of restarts within the same job | +| `num_job_retries_on_failure` | `int` | `2` | Maximum number of new job launches on failure | +| `initial_rank_heartbeat_timeout` | `int` | `1800` | Timeout for initial heartbeat (seconds) | +| `rank_heartbeat_timeout` | `int` | `300` | Timeout for subsequent heartbeats (seconds) | + +### What to Expect + +When fault tolerance is enabled and a hang or fault is detected, you should see log messages similar to: + +``` +[WARNING] [RankMonitorServer:34] Did not get subsequent heartbeat. Waited 171.92 seconds. +[WARNING] [RankMonitorServer:58] Did not get subsequent heartbeat. Waited 171.92 seconds. +FT: Simulating fault: rank_killed; rank to fail: 2 +torch.distributed.elastic.multiprocessing.api: [WARNING] Sending process 453152 closing signal SIGTERM +``` + +The system will then automatically restart training from the most recent checkpoint. + +### How It Works + +The fault tolerance system integrates with Megatron Bridge's training pipeline through several key points: + +1. **Setup Phase**: Initializes fault tolerance monitoring before training begins. +2. **Training Steps**: Wraps each training iteration with timeout monitoring. +3. **Evaluation Steps**: Monitors evaluation iterations separately. +4. **Checkpointing**: Tracks checkpoint saving operations with dedicated timeouts. +5. **State Persistence**: Saves timeout calculations to `ft_state.json` for future runs. + +The system uses a section-based approach with different timeout thresholds: +- **Setup Section**: Covers initialization and checkpoint loading. +- **Step Section**: Monitors individual training/evaluation iterations. +- **Checkpointing Section**: Tracks checkpoint saving operations. +- **Out-of-Section**: Handles time between sections. + +### Best Practices + +1. **Enable Automatic Timeout Calculation**: Set `calc_ft_timeouts=True` to let the system learn optimal timeouts from your workload. +2. **Conservative Restart Limits**: Use reasonable limits for `num_in_job_restarts` and `num_job_retries_on_failure` to avoid infinite restart loops. +3. **Monitor Logs**: Watch for fault tolerance messages to understand when and why restarts occur. +4. **Test with Simulation**: Use the fault simulation features to test your fault tolerance setup before production runs. +5. **Checkpoint Frequency**: Ensure regular checkpointing to minimize lost work during restarts. + +### Limitations + +- Currently only supported on Slurm-based clusters. +- Not compatible with NSys profiling (the plugin will automatically disable nsys if enabled). +- Checkpoint save directory must be configured and accessible. + +## Straggler Detection + +The straggler detection feature identifies slow-performing ranks and can optionally terminate training if performance falls below specified thresholds. This helps ensure efficient training by detecting and mitigating the impact of underperforming nodes. + +### Key Features + +- **Performance Monitoring**: Tracks individual and relative GPU performance scores. +- **Automatic Detection**: Identifies stragglers based on configurable thresholds. +- **Detailed Reporting**: Provides comprehensive performance reports with best/worst performing ranks. +- **Optional Termination**: Can automatically stop training when stragglers are detected. +- **Flexible Configuration**: Supports various reporting intervals and threshold settings. + +### Configuration + +Enable straggler detection through the [`bridge.training.config.NVRxStragglerDetectionConfig`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.NVRxStragglerDetectionConfig.html): + +```python +from megatron.bridge.training.config import NVRxStragglerDetectionConfig + +# Configure straggler detection in your config +config.nvrx_straggler = NVRxStragglerDetectionConfig( + enabled=True, + report_time_interval=300.0, # Report every 5 minutes + calc_relative_gpu_perf=True, + calc_individual_gpu_perf=True, + num_gpu_perf_scores_to_print=5, + gpu_relative_perf_threshold=0.7, + gpu_individual_perf_threshold=0.7, + stop_if_detected=False, # Set to True to stop training on detection + enable_logging=True, +) +``` + +### Configuration Options + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `enabled` | `bool` | `False` | Enable NVRx straggler detection | +| `report_time_interval` | `float` | `300.0` | Interval in seconds between straggler checks | +| `calc_relative_gpu_perf` | `bool` | `True` | Calculate relative GPU performance scores | +| `calc_individual_gpu_perf` | `bool` | `True` | Calculate individual GPU performance scores | +| `num_gpu_perf_scores_to_print` | `int` | `5` | Number of best/worst scores to print (0 disables periodic printing) | +| `gpu_relative_perf_threshold` | `float` | `0.7` | Threshold for relative performance (0.0-1.0) | +| `gpu_individual_perf_threshold` | `float` | `0.7` | Threshold for individual performance (0.0-1.0) | +| `stop_if_detected` | `bool` | `False` | Terminate training if stragglers are detected (saves checkpoint before exiting) | +| `enable_logging` | `bool` | `True` | Log GPU performance scores as structured data | +| `profiling_interval` | `int` | `1` | Profiling interval for the detector | +| `logger_name` | `str` | `"megatron.bridge.NVRxStragglerDetection"` | Logger name for messages | + +### Expected Output + +When straggler detection is enabled, you'll see performance reports in the training logs similar to: + +``` +GPU relative performance: + Worst performing 5/512 ranks: + Rank=76 Node=h100-001-253-012 Score=0.94 + Rank=13 Node=h100-001-010-003 Score=0.94 + Rank=45 Node=h100-001-172-026 Score=0.94 + Rank=433 Node=h100-004-141-026 Score=0.95 + Rank=308 Node=h100-003-263-012 Score=0.95 + Best performing 5/512 ranks: + Rank=432 Node=h100-004-141-026 Score=0.99 + Rank=376 Node=h100-004-005-003 Score=0.98 + Rank=487 Node=h100-004-255-026 Score=0.98 + Rank=369 Node=h100-004-004-033 Score=0.98 + Rank=361 Node=h100-004-004-023 Score=0.98 + +GPU individual performance: + Worst performing 5/512 ranks: + Rank=76 Node=h100-001-253-012 Score=0.98 + Rank=162 Node=h100-002-042-026 Score=0.98 + Rank=79 Node=h100-001-253-012 Score=0.98 + Rank=357 Node=h100-004-004-013 Score=0.98 + Rank=85 Node=h100-001-253-026 Score=0.98 + Best performing 5/512 ranks: + Rank=297 Node=h100-003-095-026 Score=1.00 + Rank=123 Node=h100-001-273-026 Score=1.00 + Rank=21 Node=h100-001-010-013 Score=1.00 + Rank=389 Node=h100-004-074-012 Score=1.00 + Rank=489 Node=h100-004-269-026 Score=1.00 + + Straggler report processing time: 0.042 sec. +``` + +If stragglers are detected and thresholds are exceeded, you'll see warnings like: + +``` +STRAGGLER DETECTION WARNING: Some GPUs have worse relative performance. Affected ranks: [76, 13, 45] +STRAGGLER DETECTION WARNING: Some GPUs performance dropped. Affected ranks: [162, 79, 357] +``` + +### Performance Scores + +The system calculates two types of performance scores: + +1. **Relative Performance**: Compares each rank's performance relative to other ranks in the same training run. +2. **Individual Performance**: Tracks each rank's performance over time to detect degradation. + +Scores range from 0.0 to 1.0, where: +- **1.0**: Best possible performance +- **0.7** (default threshold): Below this indicates a potential straggler +- **Lower values**: Indicate worse performance + +### How It Works + +The straggler detection system: + +1. **Initialization**: Sets up the NVRx detector during training setup. +2. **Monitoring**: Wraps the training step function to monitor execution time. +3. **Periodic Reporting**: Generates performance reports at specified intervals. +4. **Straggler Identification**: Compares performance scores against thresholds. +5. **Action**: Optionally saves a checkpoint and terminates training if stragglers are detected. + +### Best Practices + +1. **Appropriate Intervals**: Set `report_time_interval` based on your training characteristics. +2. **Threshold Tuning**: Adjust thresholds based on your hardware and expected performance variability. +3. **Gradual Rollout**: Start with `stop_if_detected=False` to observe performance patterns before enabling automatic termination. +4. **Monitor Logs**: Regularly check straggler reports to identify persistent hardware issues. +5. **Performance Impact**: The overhead is minimal, but you can adjust `profiling_interval` if needed. + +### Integration with Training + +The straggler detection integrates directly with the training loop: + +- Automatically initializes when [`bridge.training.resiliency.NVRxStragglerDetectionManager`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.resiliency.NVRxStragglerDetectionManager.html) is configured. +- Monitors training steps without affecting the training logic. +- Provides exit conditions that the training loop respects. +- Safely shuts down when training completes. + +## Preemption + +Training foundation models can take several hours or even days to complete. In some cases, training jobs must be halted preemptively due to cluster time limits, higher priority jobs, or other reasons. + +Megatron Bridge provides functionality to gracefully perform preemptive shutdown of training. This feature listens for user-specified signals and saves a checkpoint before exiting when the signal is received. + +### Key Features + +- **Signal-based Shutdown**: Listens for signals (default: SIGTERM) during training. +- **Graceful Exit**: Saves checkpoint before terminating to preserve training progress. +- **Distributed Coordination**: Ensures all ranks receive and handle the signal properly. +- **Flexible Configuration**: Supports different signals and timing configurations. + +### Usage Options + +Megatron Bridge provides two ways to enable preemption handling: + +#### Option 1: NeMo Run Plugin (Recommended) + +> **Warning**: This plugin is currently only supported on Slurm-based clusters. + +If you're using NeMo Run, the [`bridge.recipes.run_plugins.PreemptionPlugin`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.recipes.run_plugins.PreemptionPlugin.html) provides the simplest integration: + +```python +from megatron.bridge.recipes.run_plugins import PreemptionPlugin +import nemo_run as run + +# Configure your task +task = run.Script(...) + +# Add preemption plugin +run_plugins = [ + PreemptionPlugin( + preempt_time=60, # Send signal 60 seconds before time limit + enable_exit_handler=True, + enable_exit_handler_for_data_loader=False, + ) +] + +# Run with preemption support +run.run(task, plugins=run_plugins, executor=executor) +``` + +#### Option 2: Direct Configuration + +Configure preemption handling directly in your training configuration: + +```python +from megatron.bridge.training.config import TrainingConfig +import signal + +# Configure preemption in training config +config.train = TrainingConfig( + exit_signal_handler=True, + exit_signal=signal.SIGTERM, # Signal to listen for + exit_signal_handler_for_dataloader=False, + # ... other training config options +) +``` + +### Configuration Options + +#### PreemptionPlugin Options + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `preempt_time` | `int` | `60` | Time in seconds before job limit to send preemption signal | +| `enable_exit_handler` | `bool` | `True` | Enable the exit signal handler in training | +| `enable_exit_handler_for_data_loader` | `bool` | `False` | Enable signal handler for dataloader workers | + +#### Training Configuration Options + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `exit_signal_handler` | `bool` | `False` | Enable signal handler for graceful shutdown | +| `exit_signal` | `int` | `signal.SIGTERM` | Signal to listen for (default: SIGTERM) | +| `exit_signal_handler_for_dataloader` | `bool` | `False` | Enable signal handler for dataloader workers | + +### Expected Behavior + +When a preemption signal is received, you'll see log messages similar to: + +``` +Received signal 15, initiating graceful stop +Signal handler installed for 15 +exiting program after receiving SIGTERM. +``` + +The system will: +1. **Detect the signal** at the end of the current training step. +2. **Save a checkpoint** to preserve training progress. +3. **Log the shutdown reason** for debugging purposes. +4. **Exit gracefully** with proper cleanup. + +### How It Works + +The preemption system operates through several components: + +1. **Signal Handler Installation**: Sets up a distributed signal handler using [`bridge.training.resiliency.DistributedSignalHandler`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.resiliency.DistributedSignalHandler.html). +2. **Signal Detection**: Checks for received signals at the end of each training step. +3. **Distributed Coordination**: Uses all-gather to ensure all ranks are aware of the signal. +4. **Checkpoint Saving**: Automatically saves a checkpoint before exiting. +5. **Graceful Shutdown**: Properly cleans up resources and exits. + +### Signal Handling Details + +The `DistributedSignalHandler` class provides: +- **Cross-rank coordination**: Ensures all ranks handle the signal consistently. +- **Original handler preservation**: Restores original signal handlers on exit. +- **Flexible signal support**: Can handle different signal types (SIGTERM, SIGINT, etc.). + +### Integration with Slurm + +When using Slurm, the system automatically: +- **Receives SIGTERM** when approaching job time limits. +- **Coordinates across nodes** to ensure consistent shutdown. +- **Saves progress** before the job is forcibly terminated. + +### Best Practices + +1. **Use Appropriate Timing**: Set `preempt_time` to allow sufficient time for checkpoint saving. +2. **Monitor Logs**: Watch for preemption messages to understand shutdown patterns. +3. **Test Signal Handling**: Verify preemption works correctly in your environment. +4. **Regular Checkpointing**: Ensure regular checkpoint intervals to minimize potential data loss. +5. **Resource Cleanup**: The system handles cleanup automatically, but monitor for any resource leaks. + +## Re-run State Machine + +The re-run state machine is an experimental feature that helps with attribution of unexpected results such as NaN values, spiky loss, or other computational anomalies. It works by re-running computations to determine whether issues are transient errors, persistent hardware faults, or actually correct results. + +> **Disclaimer**: This is an experimental alpha-level feature for result attribution. Nodes flagged by this system should be subjected to standard diagnostic test suites for confirmation. + +### Key Features + +- **Automatic Re-run Logic**: Detects unexpected results and automatically re-runs computations to verify reproducibility. +- **Error Attribution**: Classifies issues as transient errors, persistent errors, or correct results. +- **Multi-stage Validation**: Uses in-place re-runs and checkpoint-based re-runs on different hardware. +- **Determinism Tracking**: Can report statistics on computational non-determinism. +- **State Management**: Handles RNG state and data iterator state for reproducible re-runs. + +### How It Works + +The re-run state machine operates through several stages: + +1. **Initial Run**: Executes the training step normally, validating results. +2. **First Re-run (In-place)**: If validation fails, re-runs on the same GPU to check reproducibility. +3. **Second Re-run (Different GPU)**: If the issue is reproducible, saves checkpoint and re-runs on different hardware. +4. **Attribution**: Determines if the issue is a transient error, persistent error, or correct result. + +### Configuration + +Configure the re-run state machine through [`bridge.training.config.RerunStateMachineConfig`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.RerunStateMachineConfig.html): + +```python +from megatron.bridge.training.config import RerunStateMachineConfig + +# Configure re-run state machine in your config +config.rerun_state_machine = RerunStateMachineConfig( + rerun_mode="validate_results", # or "report_determinism_stats" or "disabled" + check_for_nan_in_loss=True, + check_for_spiky_loss=False, + spiky_loss_factor=10.0, # Adjust for your model architecture + error_injection_rate=0, # For testing only + error_injection_type="transient_error", +) +``` + +### Configuration Options + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `rerun_mode` | `str` | `"disabled"` | Operating mode: `"disabled"`, `"validate_results"`, or `"report_determinism_stats"` | +| `check_for_nan_in_loss` | `bool` | `True` | Check for NaN values in loss | +| `check_for_spiky_loss` | `bool` | `False` | Check for unexpectedly large loss values | +| `spiky_loss_factor` | `float` | `10.0` | Factor for spiky loss detection. Loss is flagged if it exceeds this multiple of max observed loss. Larger models may need higher values (e.g., 15-20 for 70B+). | +| `error_injection_rate` | `int` | `0` | Rate for injecting test errors (testing only) | +| `error_injection_type` | `str` | `"transient_error"` | Type of error to inject for testing | + +### Operating Modes + +#### 1. Disabled Mode (`disabled`) +- **Purpose**: No result validation or re-run logic. +- **Behavior**: Training proceeds normally without any result checking. +- **Use Case**: When re-run overhead is not acceptable or validation is not needed. + +#### 2. Report Stats Mode (`report_determinism_stats`) +- **Purpose**: Collect statistics on computational determinism. +- **Behavior**: Re-runs every step once to measure variability. +- **Output**: Reports on computational non-determinism without stopping training. + +#### 3. Validate Results Mode (`validate_results`) +- **Purpose**: Full validation with re-runs and hardware fault attribution. +- **Behavior**: Re-runs computations when unexpected results are detected. +- **Exit Conditions**: May exit with specific codes for checkpointing or validation failure. + +### Integration with Training + +The re-run state machine integrates at the training step level: + +```python +# In train_step function +rerun_state_machine = get_rerun_state_machine() +while rerun_state_machine.should_run_forward_backward(data_iterator): + # Execute forward-backward pass + loss_dict = forward_backward_func(...) + + # Validate results (automatically handled in forward_step) + # check_for_nan_in_loss and check_for_spiky_loss are passed to loss function + +should_checkpoint, should_exit, exit_code = rerun_state_machine.should_checkpoint_and_exit() +if should_checkpoint: + save_checkpoint(...) +if should_exit: + sys.exit(exit_code) +``` + +### Exit Codes + +The re-run state machine uses specific exit codes to control job behavior: + +- **Exit Code 16** (`EXIT_CODE_RESUME_TO_DISAMBIGUATE`): Job should be restarted from checkpoint to re-run on different hardware. +- **Exit Code 17** (`EXIT_CODE_FAILED_ON_RESULT_VALIDATION`): Job failed validation and should not continue. + +### Expected Behavior + +#### Validation Success +When validation passes, training continues normally with no additional overhead. + +#### Transient Error Detection +``` +Unexpected result tensor(nan) on rank 0 at iteration #150 invocation #1 (message='loss is NaN') +First rerun: unexpected result is not reproducible within the tolerance +Possible transient error! +``` + +#### Persistent Error Detection +``` +First rerun: unexpected result is reproducible within the tolerance +Need to rerun on a different GPU to verify correctness +Second rerun: unexpected result is not reproducible on a different GPU, therefore was likely incorrect +Possible persistent error! +``` + +#### Correct Result (False Positive) +``` +Second rerun: unexpected result is reproducible on a different GPU, therefore it was likely correct +Correct result (but possible Application error) +``` + +### Result Attribution Categories + +1. **Transient Error**: Result not reproducible on same GPU - likely temporary hardware glitch. +2. **Persistent Error**: Result reproducible on same GPU but different on other GPU - likely hardware fault. +3. **Correct Result**: Result reproducible across different GPUs - likely correct but unexpected. + +### Data Iterator Integration + +The system uses `RerunDataIterator` to handle data replay: +- **State Saving**: Captures data iterator state for reproducible re-runs. +- **Replay Capability**: Can rewind and replay the same data batches. +- **Checkpoint Support**: Saves/restores iterator state across job restarts. + +## In-Process Restart + +> **Warning**: This is a highly experimental feature and is subject to change in backwards incompatible ways without notice. + +The in-process restart mechanism provides automatic fault recovery by restarting the training function within the same operating system process when failures occur. Unlike traditional scheduler-level restarts, in-process restart eliminates the overhead of launching new jobs, starting containers, initializing Python interpreters, and creating new CUDA contexts. + +> **Note**: In-process restart is not suitable for all types of failures. Hardware-level failures such as switch failures, network partitions, or multiple node failures that render nodes inaccessible cannot be recovered through in-process restart alone. For comprehensive fault tolerance, it is recommended to combine in-process restart with the fault tolerance system (in-job restarts) described earlier in this document. This layered approach provides both fast recovery for software faults and robust handling of hardware-level failures. + +For comprehensive information about this functionality, refer to the [NVIDIA Resiliency Extension In-Process Restart documentation](https://nvidia.github.io/nvidia-resiliency-ext/inprocess/index.html). + +### Key Features + +- **In-Process Recovery**: Restarts training within the same process, avoiding container and interpreter restart overhead. +- **Automatic Fault Detection**: Detects unhandled Python exceptions, deadlocks, and livelocks across all distributed ranks. +- **Coordinated Restart**: Ensures all healthy ranks restart simultaneously when any rank encounters a fault. +- **Timeout Mechanisms**: Provides both soft and hard timeouts to detect and recover from hangs. +- **Rank Reassignment**: Supports excluding unhealthy ranks and utilizing warm reserve workers. +- **State Reuse**: Enables reuse of process-group-independent objects across restart attempts to minimize latency. +- **Granular Control**: Supports both node-level and rank-level restart granularity. +- **Health Checks**: Performs GPU health validation and optionally tracks fault counts. + +### Prerequisites + +Before using in-process restart, ensure the following requirements are met: + +1. **PyTorch Version**: PyTorch v2.5.1 or higher is required. +2. **NCCL Version**: NCCL v2.26.2 or higher is required. +3. **Checkpoint Configuration**: A valid checkpoint directory must be configured for state recovery. +4. **GIL-Released Operations**: All operations that wait on NCCL kernels or synchronize with GPU must release the Python Global Interpreter Lock (GIL). + +> **Important**: If operations hold the GIL during a fault, graceful restart cannot proceed, and affected ranks will be forcibly terminated. + +### Configuration + +Configure in-process restart through [`bridge.training.config.InProcessRestartConfig`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.InProcessRestartConfig.html): + +```python +from megatron.bridge.training.config import InProcessRestartConfig + +# Configure in-process restart in your config +config.inprocess_restart = InProcessRestartConfig( + enabled=True, + active_world_size=None, # Defaults to WORLD_SIZE, set lower to use warm reserves + granularity="node", # or "rank" for rank-level restart + max_iterations=None, # No limit on restart attempts + soft_timeout=60.0, # Timeout for detecting GIL-released hangs + hard_timeout=90.0, # Timeout for forcibly terminating hung ranks + heartbeat_interval=30.0, + heartbeat_timeout=60.0, + monitor_thread_interval=1.0, + monitor_process_interval=1.0, + progress_watchdog_interval=1.0, + barrier_timeout=120.0, + completion_timeout=120.0, + last_call_wait=1.0, + termination_grace_time=1.0, + empty_cuda_cache=True, + max_rank_faults=None, # No limit on rank faults + monitor_process_logdir=None, # Disable monitor process logging +) +``` + +### Configuration Options + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `enabled` | `bool` | `False` | Enable in-process restart mechanism | +| `active_world_size` | `int` | `None` | Number of ranks initially executing workload (remaining ranks are warm reserves) | +| `granularity` | `str` | `"node"` | Restart granularity: `"node"` or `"rank"` | +| `max_iterations` | `int` | `None` | Maximum number of restart iterations (None = unlimited) | +| `soft_timeout` | `float` | `60.0` | Soft progress timeout in seconds (for detecting GIL-released hangs) | +| `hard_timeout` | `float` | `90.0` | Hard progress timeout in seconds (for forcibly terminating hung ranks) | +| `heartbeat_interval` | `float` | `30.0` | Interval in seconds for heartbeat monitoring | +| `heartbeat_timeout` | `float` | `60.0` | Timeout in seconds for detecting missing rank heartbeats | +| `monitor_thread_interval` | `float` | `1.0` | Monitoring interval in seconds for monitoring thread | +| `monitor_process_interval` | `float` | `1.0` | Monitoring interval in seconds for monitoring process | +| `progress_watchdog_interval` | `float` | `1.0` | Interval in seconds for automatic progress watchdog updates | +| `barrier_timeout` | `float` | `120.0` | Timeout in seconds for internal distributed barriers | +| `completion_timeout` | `float` | `120.0` | Timeout in seconds for completion barrier on all ranks | +| `last_call_wait` | `float` | `1.0` | Time interval in seconds for other ranks to report concurrent failures | +| `termination_grace_time` | `float` | `1.0` | Interval in seconds between SIGTERM and SIGKILL on hard timeout | +| `empty_cuda_cache` | `bool` | `True` | Empty CUDA cache during restart finalization | +| `max_rank_faults` | `int` | `None` | Maximum number of rank faults allowed before terminating (None = unlimited) | +| `monitor_process_logdir` | `str` | `None` | Directory for monitor process log files (None = disabled) | + +### Slurm Configuration Requirements + +> **Warning**: Running in-process restart through NeMo-Run's Slurm Executor is **not currently supported**. + +If you need to use in-process restart with Slurm, you must launch your jobs directly using `srun` with the proper configuration. Refer to the [NVIDIA Resiliency Extension Slurm configuration guide](https://nvidia.github.io/nvidia-resiliency-ext/inprocess/usage_guide.html#running-with-slurm) for detailed instructions on: + +- Setting `--kill-on-bad-exit=0` to prevent Slurm from terminating the entire job on rank failures +- Using the `wait_daemon.py` utility for proper monitoring process cleanup +- Configuring SLURM PMI for compatibility + +#### Monitor Process Log Files + +When `monitor_process_logdir` is configured, the system automatically generates monitor process log files for rank 0 only. The log file path must be coordinated between your Python configuration and the `wait_daemon.py` script used in your Slurm launch command. + +The system creates log files with the following naming convention: + +``` +monitor_{SLURM_JOB_ID}_{hostname}_{SLURM_PROCID}_{SLURM_LOCALID}.log +``` + +Where: +- `SLURM_JOB_ID`: The Slurm job ID from the `SLURM_JOB_ID` environment variable +- `hostname`: The hostname of the node where rank 0 is running +- `SLURM_PROCID`: The global rank from the `SLURM_PROCID` environment variable +- `SLURM_LOCALID`: The local rank on the node from the `SLURM_LOCALID` environment variable + +**Python Configuration:** + +```python +config.inprocess_restart = InProcessRestartConfig( + enabled=True, + monitor_process_logdir="/scratch/logs/monitor", # Provide directory only +) +``` + +**Corresponding Slurm Launch Command:** + +You must pass the same log file path pattern to `wait_daemon.py` in your sbatch script. The path should include `{rank}` as a placeholder that will be substituted with the actual rank: + +```bash +srun --kill-on-bad-exit=0 \ + python -m nvidia_resiliency_ext.inprocess.wait_daemon \ + --monitor-process-logfile=/scratch/logs/monitor/monitor_${SLURM_JOB_ID}_$(hostname)_\${SLURM_PROCID}_\${SLURM_LOCALID}.log \ + -- \ + python your_training_script.py +``` + +> **Important**: The monitor process log file path must match between your Python configuration (`monitor_process_logdir`) and the `wait_daemon.py` command-line argument. This coordination ensures that `wait_daemon.py` can properly monitor and wait for the monitor process to complete its cleanup before exiting. + +### Integration in Megatron Bridge + +The in-process restart system integrates with Megatron Bridge's training pipeline through several mechanisms: + +#### 1. Function Wrapping + +The `pretrain()` function detects when in-process restart is enabled and wraps the internal `_pretrain()` function with the restart mechanism: + +```python +if config.inprocess_restart and config.inprocess_restart.enabled: + from megatron.bridge.training.inprocess_restart import maybe_wrap_for_inprocess_restart + + wrapped_pretrain, store = maybe_wrap_for_inprocess_restart( + _pretrain, config.inprocess_restart, state + ) + wrapped_pretrain(state, forward_step_func, store=store) +``` + +#### 2. Coordination Store + +A dedicated `TCPStore` is created for coordination between ranks during restart operations: +- Uses `MASTER_PORT + 1` to avoid conflicts with PyTorch distributed +- Enables rank-to-rank communication for fault detection and recovery +- Supports prefix-based isolation for each restart attempt + +#### 3. State Management + +During restart, the system performs comprehensive cleanup: + +- **PyTorch State**: Destroys distributed process groups via `torch.distributed.destroy_process_group()` +- **Megatron State**: Cleans up global state through `destroy_global_state()` +- **Training State**: Resets the `GlobalState` object for fresh initialization +- **CUDA State**: Optionally empties CUDA cache to free memory +- **Async Workers**: Aborts persistent async checkpoint worker processes + +#### 4. Restart Flow + +When a fault occurs on any rank: + +1. **Fault Detection**: The wrapper detects the exception, timeout, or missing heartbeat +2. **Distributed Abort**: All ranks are notified and begin coordinated shutdown +3. **State Cleanup**: Each rank cleans up PyTorch, Megatron, and CUDA state +4. **Health Check**: GPU health is validated on each rank +5. **Rank Reassignment**: Unhealthy ranks are excluded, reserves may be activated +6. **Barrier Synchronization**: All healthy ranks wait at a distributed barrier +7. **Function Restart**: The wrapped function restarts on all healthy ranks simultaneously + +#### 5. Restart Iterations + +The `CallWrapper` tracks restart iterations and provides this information to the wrapped function: +- Iteration 0: Initial execution +- Iteration 1+: Restart attempts after faults +- Used to create isolated `PrefixStore` instances per restart attempt + +### Environment Configuration + +#### Required Environment Variables + +Set these environment variables to optimize in-process restart behavior: + +```bash +# Suppress c10d TCPStore wait timeout warnings +export TORCH_CPP_LOG_LEVEL=error + +# Prevent PyTorch NCCL Watchdog from forcibly terminating on NCCL/CUDA errors +export TORCH_NCCL_RETHROW_CUDA_ERRORS=0 + +# Disable NVLS support in NCCL (required for in-process restart) +export NCCL_NVLS_ENABLE=0 +``` + +#### PyTorch NCCL Watchdog Timeout + +Configure the PyTorch NCCL watchdog timeout to be longer than the `hard_timeout`: + +```python +import torch.distributed as dist +from datetime import timedelta + +# When initializing the distributed backend +dist.init_process_group( + backend='nccl', + timeout=timedelta(seconds=config.inprocess_restart.hard_timeout + 60) +) +``` + +### Known Issues + +Refer to the [NVIDIA Resiliency Extension Known Issues](https://nvidia.github.io/nvidia-resiliency-ext/inprocess/usage_guide.html#known-issues) for the most up-to-date list of limitations and workarounds related to: + +- PyTorch distributed limitations +- NCCL collective termination +- CUDA context handling +- Checkpoint worker cleanup diff --git a/fern/v0.2.0/pages/training/training-loop-settings.mdx b/fern/v0.2.0/pages/training/training-loop-settings.mdx new file mode 100644 index 0000000000..3909cb2827 --- /dev/null +++ b/fern/v0.2.0/pages/training/training-loop-settings.mdx @@ -0,0 +1,88 @@ +--- +title: "Training Loop Configuration" +description: "" +--- + +The [`bridge.training.config.TrainingConfig`](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.training.config.TrainingConfig.html) contains settings related to the training loop bounds, exit conditions, validation, batch sizing, and memory management. + +## Key Parameters + +Configure these parameters to control core training behavior, resource utilization, and monitoring across distributed setups. + +### Batch Configuration +Define how data is batched and distributed across devices during training. + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `micro_batch_size` | `Optional[int]` | `None` | Batch size per model instance (local batch size) | +| `global_batch_size` | `Optional[int]` | `None` | Training batch size across all devices | +| `rampup_batch_size` | `Optional[list[int]]` | `None` | Batch size ramp up: `[start_size, increment, ramp_samples]` | +| `decrease_batch_size_if_needed` | `bool` | `False` | Automatically decrease batch size if needed for fault tolerance | + +The relationship between batch sizes: +- **Global batch size** = `micro_batch_size` × `data_parallel_size` × `gradient_accumulation_steps` +- If `global_batch_size` is not set, it defaults to `micro_batch_size` × `data_parallel_size` + +### Training Duration + +Control when training stops using iteration counts, sample counts, or time-based limits. + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `train_iters` | `Optional[int]` | `None` | Total number of iterations to train | +| `train_samples` | `Optional[int]` | `None` | Total number of samples to train | +| `exit_interval` | `Optional[int]` | `None` | Exit after iteration divisible by this value | +| `exit_duration_in_mins` | `Optional[int]` | `None` | Exit after this many minutes | + +**Training Mode Selection** + +Megatron-Bridge supports two modes for specifying training duration: + +1. **Iteration-based training**: Specify `train_iters` to control the total number of training iterations. +2. **Sample-based training**: Specify `train_samples` to control the total number of training samples. + +**Important constraints:** +- You must specify **exactly one** of `train_iters` or `train_samples` - not both. +- When using `train_samples`, training iterations are automatically calculated as `train_samples // global_batch_size`. +- Batch size rampup (`rampup_batch_size`) is not currently supported with sample-based training. +- Your scheduler configuration should match your training mode (see [Learning Rate Scheduling](/optimizer-scheduler#learning-rate-scheduling)). + +### Validation +Configure validation frequency, duration, and evaluation-only modes. + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `eval_iters` | `int` | `100` | Number of iterations for validation/test evaluation | +| `eval_interval` | `Optional[int]` | `1000` | Interval between validation runs | +| `skip_train` | `bool` | `False` | Skip training, only do evaluation and exit | + +**Note:** To control validation behavior: +- Set `eval_iters` to `0` to disable validation entirely (both during and after training). +- Set `eval_interval` to `None` to skip validation during training, but still run validation after training completes. + +### Memory Management +Control GPU memory cleanup and garbage collection to prevent memory issues during training. + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `empty_unused_memory_level` | `Literal[0, 1, 2]` | `0` | Call `torch.cuda.empty_cache()` each iteration (0=off, 1=moderate, 2=aggressive) | +| `manual_gc` | `bool` | `False` | Synchronize Python garbage collection across ranks to avoid stragglers | +| `manual_gc_interval` | `int` | `0` | Training step interval for manual garbage collection (0=disabled) | +| `manual_gc_eval` | `bool` | `True` | Enable garbage collection during evaluation when using manual GC | + +### Signal Handling and Exit Conditions +Set up automatic checkpoint saving and clean exit procedures for signal-based interruptions. + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `exit_signal_handler` | `bool` | `False` | Save checkpoint and shutdown gracefully on signal detection | +| `exit_signal` | `int` | `signal.SIGTERM` | Signal to handle for graceful shutdown | +| `exit_signal_handler_for_dataloader` | `bool` | `False` | Use signal handler for dataloader workers | + +### Performance Monitoring +Monitor training consistency and synchronization across distributed processes. + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `check_weight_hash_across_dp_replicas_interval` | `Optional[int]` | `None` | Check weight hash consistency across data parallel replicas | +| `train_sync_interval` | `Optional[int]` | `None` | CPU-GPU synchronization interval to prevent CPU running ahead | diff --git a/fern/versions/v0.2.0.yml b/fern/versions/v0.2.0.yml new file mode 100644 index 0000000000..0daca9428b --- /dev/null +++ b/fern/versions/v0.2.0.yml @@ -0,0 +1,141 @@ +navigation: + - section: Home + contents: + - page: Welcome + path: ../v0.2.0/pages/index.mdx + - section: Guides + contents: + - page: Parallelisms + path: ../v0.2.0/pages/parallelisms.mdx + - page: Performance Summary + path: ../v0.2.0/pages/performance-summary.mdx + - page: Performance Guide + path: ../v0.2.0/pages/performance-guide.mdx + - page: Recipe Usage + path: ../v0.2.0/pages/recipe-usage.mdx + - page: NeMo 2 Migration Guide + path: ../v0.2.0/pages/nemo2-migration-guide.mdx + - page: Megatron-LM to Megatron-Bridge + path: ../v0.2.0/pages/megatron-lm-to-megatron-bridge.mdx + - section: Bridge with Hugging Face + contents: + - page: Bridge Guide + path: ../v0.2.0/pages/bridge-guide.mdx + - page: Bridge Tech Details + path: ../v0.2.0/pages/bridge-tech-details.mdx + - section: Models + contents: + - page: Large Language Models + path: ../v0.2.0/pages/models/llm/index.mdx + - page: DeepSeek V2 + path: ../v0.2.0/pages/models/llm/deepseek-v2.mdx + - page: DeepSeek V3 + path: ../v0.2.0/pages/models/llm/deepseek-v3.mdx + - page: Gemma 2 + path: ../v0.2.0/pages/models/llm/gemma2.mdx + - page: Gemma 3 + path: ../v0.2.0/pages/models/llm/gemma3.mdx + - page: GLM 4.5 + path: ../v0.2.0/pages/models/llm/glm45.mdx + - page: GPT-OSS + path: ../v0.2.0/pages/models/llm/gpt-oss.mdx + - page: LLaMA 3 + path: ../v0.2.0/pages/models/llm/llama3.mdx + - page: LLaMA Nemotron + path: ../v0.2.0/pages/models/llm/llama-nemotron.mdx + - page: Mistral + path: ../v0.2.0/pages/models/llm/mistral.mdx + - page: Moonlight + path: ../v0.2.0/pages/models/llm/moonlight.mdx + - page: Nemotron 3 + path: ../v0.2.0/pages/models/llm/nemotron3.mdx + - page: Nemotron-H + path: ../v0.2.0/pages/models/llm/nemotronh.mdx + - page: OlMoE + path: ../v0.2.0/pages/models/llm/olmoe.mdx + - page: Qwen + path: ../v0.2.0/pages/models/llm/qwen.mdx + - page: Vision Language Models + path: ../v0.2.0/pages/models/vlm/index.mdx + - page: Qwen 2.5 VL + path: ../v0.2.0/pages/models/vlm/qwen2.5-vl.mdx + - page: Qwen 3 VL + path: ../v0.2.0/pages/models/vlm/qwen3-vl.mdx + - page: Nemotron Nano V2 VL + path: ../v0.2.0/pages/models/vlm/nemotron-nano-v2-vl.mdx + - page: Ministral 3 + path: ../v0.2.0/pages/models/vlm/ministral3.mdx + - page: GLM 4.5V + path: ../v0.2.0/pages/models/vlm/glm-45v.mdx + - page: Gemma 3 VL + path: ../v0.2.0/pages/models/vlm/gemma3-vl.mdx + - section: Training + contents: + - page: Training Overview + path: ../v0.2.0/pages/training/README.mdx + - page: Config Container Overview + path: ../v0.2.0/pages/training/config-container-overview.mdx + - page: Entry Points + path: ../v0.2.0/pages/training/entry-points.mdx + - page: Training Loop Settings + path: ../v0.2.0/pages/training/training-loop-settings.mdx + - page: Optimizer and Scheduler + path: ../v0.2.0/pages/training/optimizer-scheduler.mdx + - page: Logging + path: ../v0.2.0/pages/training/logging.mdx + - page: Profiling + path: ../v0.2.0/pages/training/profiling.mdx + - page: Checkpointing + path: ../v0.2.0/pages/training/checkpointing.mdx + - page: Resiliency + path: ../v0.2.0/pages/training/resiliency.mdx + - page: Mixed Precision + path: ../v0.2.0/pages/training/mixed-precision.mdx + - page: Communication Overlap + path: ../v0.2.0/pages/training/communication-overlap.mdx + - page: Attention Optimizations + path: ../v0.2.0/pages/training/attention-optimizations.mdx + - page: Activation Recomputation + path: ../v0.2.0/pages/training/activation-recomputation.mdx + - page: CPU Offloading + path: ../v0.2.0/pages/training/cpu-offloading.mdx + - page: PEFT + path: ../v0.2.0/pages/training/peft.mdx + - page: Packed Sequences + path: ../v0.2.0/pages/training/packed-sequences.mdx + - page: Multi-Token Prediction + path: ../v0.2.0/pages/training/multi-token-prediction.mdx + - page: Distillation + path: ../v0.2.0/pages/training/distillation.mdx + - page: Pruning + path: ../v0.2.0/pages/training/pruning.mdx + - page: Callbacks + path: ../v0.2.0/pages/training/callbacks.mdx + - section: Model Optimization + contents: + - page: ModelOpt Overview + path: ../v0.2.0/pages/modelopt/README.mdx + - page: Quantization + path: ../v0.2.0/pages/modelopt/quantization.mdx + - section: Development + contents: + - page: Adding New Models + path: ../v0.2.0/pages/adding-new-models.mdx + - page: Bridge RL Integration + path: ../v0.2.0/pages/bridge-rl-integration.mdx + - page: Documentation + path: ../v0.2.0/pages/documentation.mdx + - section: API Reference + contents: + - library: megatron-bridge + title: Library Reference + - section: Releases + contents: + - page: Releases Overview + path: ../v0.2.0/pages/releases/README.mdx + - page: Software Versions + path: ../v0.2.0/pages/releases/software-versions.mdx + - page: Changelog + path: ../v0.2.0/pages/releases/changelog.mdx + - page: Known Issues + path: ../v0.2.0/pages/releases/known-issues.mdx