docs: fern migration - #2410
Closed
lbliii wants to merge 3 commits into
Closed
Conversation
Signed-off-by: Lawrence Lane <llane@nvidia.com>
Contributor
📝 WalkthroughWalkthroughThis PR establishes a comprehensive Fern documentation framework for Megatron Bridge, introducing 100+ MDX documentation pages covering models, training workflows, optimization, and APIs, alongside utility scripts for converting MyST/Sphinx documentation to Fern MDX format and organizing documentation assets. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 3 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
Tip Issue Planner is now in beta. Read the docs and try it out! Share your feedback on Discord. Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
| return True | ||
|
|
||
|
|
||
| def main() -> None: |
Contributor
There was a problem hiding this comment.
Pipeline failure: add docstring to main().
CI reports D103 Missing docstring in public function at line 42.
Proposed fix
def main() -> None:
+ """CLI entrypoint: add frontmatter to all MDX files under the given pages directory."""
parser = argparse.ArgumentParser(🧰 Tools
🪛 GitHub Actions: CICD NeMo
[error] 42-42: D103 Missing docstring in public function
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@fern/scripts/add_frontmatter.py` at line 42, Public function main() is
missing a docstring; add a clear triple-quoted docstring immediately under the
def main() line in the add_frontmatter.py script that briefly explains what
main() does, any side effects (e.g., reads/writes files or modifies
frontmatter), expected inputs or configuration, and that it returns None.
Reference the main() function and ensure the docstring follows the project's
style (one-line summary plus optional short description) to satisfy D103.
| SKIP_DIRS = {"_templates", "_build", "apidocs", ".venv", ".git"} | ||
|
|
||
|
|
||
| def main() -> None: |
Contributor
There was a problem hiding this comment.
Pipeline failure: add docstring to main().
The CI pipeline reports D103 Missing docstring in public function at line 22. Add a Google-style docstring to fix it.
Proposed fix
def main() -> None:
+ """Copy docs/*.md into the Fern versioned pages directory as .mdx files."""
parser = argparse.ArgumentParser(📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| def main() -> None: | |
| def main() -> None: | |
| """Copy docs/*.md into the Fern versioned pages directory as .mdx files.""" |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@fern/scripts/copy_docs_to_fern.py` at line 21, Add a Google-style docstring
to the public function main() describing its purpose, any arguments (none),
return value (None), and side effects (e.g., copies documentation files or
invokes the script's workflow) so the D103 lint error is resolved; update the
docstring inside the main() function definition (def main() -> None:) and keep
it short, using the Google style "Args:", "Returns:", and an optional "Raises:"
section only if the function can raise exceptions.
| return False | ||
|
|
||
|
|
||
| def main() -> None: |
Contributor
There was a problem hiding this comment.
Pipeline failure: add docstring to main().
CI reports D103 Missing docstring in public function at line 54.
Proposed fix
def main() -> None:
+ """CLI entrypoint: expand {include} directives in target MDX files."""
parser = argparse.ArgumentParser(🧰 Tools
🪛 GitHub Actions: CICD NeMo
[error] 54-54: D103 Missing docstring in public function
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@fern/scripts/expand_includes.py` at line 54, Add a module docstring for the
public function main(): open the function definition for main() and add a proper
triple-quoted docstring immediately below the def that briefly describes what
main does and notes that it takes no arguments and returns None (e.g., one-line
summary plus optional short description), following D103 style conventions;
ensure the docstring uses triple quotes and is indented inside the main()
function.
| return False | ||
|
|
||
|
|
||
| def main() -> None: |
Contributor
There was a problem hiding this comment.
Missing docstring on main() — pipeline failure.
The CI pipeline reports D103 Missing docstring in public function at line 34. Add a docstring to main() to fix this.
📝 Proposed fix
def main() -> None:
+ """Remove duplicate H1 headings that match frontmatter titles in MDX files."""
parser = argparse.ArgumentParser(🧰 Tools
🪛 GitHub Actions: CICD NeMo
[error] 34-34: D103 Missing docstring in public function
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@fern/scripts/remove_duplicate_h1.py` at line 34, Add a concise docstring to
the public function main() describing its purpose and behavior to satisfy the
D103 linter error; update the main() function's definition (symbol: main) to
include a short one- or two-sentence triple-quoted docstring explaining what the
script does and any important side effects or return behavior.
Comment on lines
+211
to
+222
| # 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)) | ||
|
|
Contributor
There was a problem hiding this comment.
Lambda in forward_step_fn ignores its argument — incorrect with pipeline parallelism.
The loss-function lambda captures outputs from the enclosing scope instead of using its _out parameter. With pipeline_model_parallel_size > 1, the framework passes the actual last-stage output to the loss function; ignoring it means the loss is computed on the wrong (intermediate-stage) tensor. Since the document explicitly targets TP/PP/CP setups, this will mislead readers.
Proposed fix
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))
+ return outputs, (lambda out: rl_loss_fn(out, batch))🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@fern/v0.2.0/pages/bridge-rl-integration.mdx` around lines 211 - 222, The
lambda returned by forward_step_fn captures the outer variable outputs instead
of using the value passed by the pipeline, so with pipeline_model_parallel_size
> 1 the loss will be computed on an intermediate tensor; update the returned
collector to accept and use its argument (e.g., change (lambda _out:
rl_loss_fn(outputs, batch)) to a form that calls rl_loss_fn with the lambda
parameter such as lambda last_stage_out: rl_loss_fn(last_stage_out, batch)) so
rl_loss_fn receives the actual last-stage output; keep the function name
forward_step_fn and the batch/rl_loss_fn references to locate the change.
Comment on lines
+60
to
+64
| ## 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) |
Contributor
There was a problem hiding this comment.
Potentially broken relative links.
Line 61 uses /../llm/gemma3 (single ..) while lines 62–64 use the /../../ pattern flagged in other files. Verify these resolve correctly in Fern.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@fern/v0.2.0/pages/models/vlm/gemma3-vl.mdx` around lines 60 - 64, The
relative links under "Related Docs" may be inconsistent: the entry for Text-Only
Models uses "/../llm/gemma3" while the others use "/../../"; verify and
normalize them so they resolve correctly in Fern. Update the "/../llm/gemma3"
link to the same pattern as the other links (e.g., "/../../llm/gemma3") or
change all to the correct relative path used across the docs, ensuring the link
strings "Text-Only Models: [Gemma 3](/../llm/gemma3)" and the
recipe/config/entry-points links use a consistent and resolvable relative path.
| --hf-model-path $HF_MODEL_PATH \ | ||
| --pretrained-checkpoint $MEGATRON_MODEL_PATH \ | ||
| --lora-on-language-model \ | ||
| —-lora-on-vision-model \ |
Contributor
There was a problem hiding this comment.
Bug: em dash (—) instead of double hyphen (--) in CLI flag.
—-lora-on-vision-model starts with an em dash (U+2014) rather than two ASCII hyphens. Users who copy-paste this command will get an unrecognized-argument error.
Proposed fix
-—-lora-on-vision-model \
+--lora-on-vision-model \📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| —-lora-on-vision-model \ | |
| --lora-on-vision-model \ |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@fern/v0.2.0/pages/models/vlm/nemotron-nano-v2-vl.mdx` at line 134, Replace
the em dash in the CLI flag token `—-lora-on-vision-model` with two ASCII
hyphens so the flag reads `--lora-on-vision-model`; locate the token in the VLM
page content (the CLI example containing `—-lora-on-vision-model`) and update
the character U+2014 to two ASCII '-' characters to prevent
unrecognized-argument errors when users copy-paste the command.
Comment on lines
+310
to
+358
| ```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` |
Contributor
There was a problem hiding this comment.
Combined parallelism example has an invalid EP configuration.
With the given settings (TP=2, PP=4, CP=2) on 32 GPUs, the data-parallel size is 2 (as correctly noted on Line 358). However, expert_model_parallel_size=4 requires the DP size to be at least 4 (since EP is carved out of the DP dimension). With DP=2 and EP=4, this configuration would fail at runtime.
Consider either increasing the world size (e.g., 64 GPUs → DP=4) or reducing expert_model_parallel_size to 2.
Proposed fix (option: reduce EP to fit DP=2)
# Expert parallelism (for MoE models)
num_moe_experts=8, # 8 experts
- expert_model_parallel_size=4, # Distribute experts across 4 GPUs
- expert_tensor_parallel_size=2, # Apply TP within each expert
+ expert_model_parallel_size=2, # Distribute experts across 2 GPUs (must be <= DP size)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@fern/v0.2.0/pages/parallelisms.mdx` around lines 310 - 358, The example’s
expert parallelism (expert_model_parallel_size=4) is invalid for the computed
data-parallel size (data_parallel_size = world_size /
(tensor_model_parallel_size × pipeline_model_parallel_size ×
context_parallel_size) = 2); update the config so EP ≤ DP by either increasing
world_size (e.g., set world_size to 64 so data_parallel_size becomes 4) or
lowering expert_model_parallel_size to 2, and update the prose/example
accordingly to show consistent values for tensor_model_parallel_size,
pipeline_model_parallel_size, context_parallel_size, expert_model_parallel_size,
and the resulting data_parallel_size.
|
|
||
| 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. |
Contributor
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Check if there's a heading or anchor target matching either pattern
echo "=== Searching for anchor patterns ==="
rg -n 'bridge\.recipes\.run_plugins|bridge-recipes-run_plugins|bridge-recipes-run-plugins' --type-add 'mdx:*.mdx' --type=mdx fern/ 2>/dev/null || echo "No matches found"
echo -e "\n=== Checking the file around line 246 ==="
cat -n fern/v0.2.0/pages/recipe-usage.mdx | sed -n '240,260p'
echo -e "\n=== Searching for any heading containing 'run_plugins' or 'run plugins' ==="
rg -n 'run.plugins|run_plugins' --type-add 'mdx:*.mdx' --type=mdx fern/v0.2.0/pages/recipe-usage.mdx 2>/dev/null || echo "No matches found"Repository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 5446
🏁 Script executed:
#!/bin/bash
echo "=== Checking link patterns in recipe-usage.mdx ==="
rg -n '\[.*\]\(#' fern/v0.2.0/pages/recipe-usage.mdx | head -20
echo -e "\n=== Checking for API reference sections ==="
rg -n '## .*[Aa]pi|## .*[Rr]eference' fern/v0.2.0/pages/recipe-usage.mdx
echo -e "\n=== Looking for external API doc patterns in the directory ==="
rg -n 'apidocs|api reference' --type-add 'mdx:*.mdx' --type=mdx fern/v0.2.0/pages/ | head -10Repository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 3717
🏁 Script executed:
#!/bin/bash
echo "=== Checking all headings in recipe-usage.mdx ==="
rg -n '^#+' fern/v0.2.0/pages/recipe-usage.mdx | head -30
echo -e "\n=== Checking if there's a pattern for module-level API docs URL ==="
rg -n 'apidocs.*run_plugins|bridge\.recipes' fern/v0.2.0/pages/ --type-add 'mdx:*.mdx' --type=mdx | grep -E '(https://|apidocs)' | head -5Repository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 2664
Replace internal anchor link with external API documentation URL.
The anchor link #bridge.recipes.run_plugins has no corresponding heading in this file and will not resolve. Replace with the external API reference URL following the pattern used throughout the documentation:
[API reference](https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.recipes.run_plugins.html)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@fern/v0.2.0/pages/recipe-usage.mdx` at line 246, Replace the broken internal
anchor link "#bridge.recipes.run_plugins" in the text "See the [API
reference](`#bridge.recipes.run_plugins`) for a list of available NeMo-Run
plugins." with the external API docs URL used elsewhere:
"https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.recipes.run_plugins.html"
so the link targets the correct external API reference for
bridge.recipes.run_plugins.
Comment on lines
+81
to
+135
| <Note> | ||
| 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: | ||
|
|
||
| </Note> |
Contributor
There was a problem hiding this comment.
Broken markup: </Note> tag is misplaced and directory tree lacks a code fence.
The <Note> opened on line 81 is not closed until line 135, which means the "Checkpoint Loading Strictness" section (line 90), the "Fine-tuning Configuration" section (line 103), "Checkpoint Format" (line 109), "Performance Optimizations" (line 115), and "Checkpoint Contents" (line 123) are all inadvertently wrapped inside the <Note> block. This will cause rendering issues.
Additionally, the directory structure starting at line 136 is missing an opening code fence (```), while a closing fence exists at line 159.
Proposed fix
Close the <Note> after line 88 and add an opening code fence before the directory tree:
- `load` + `ckpt_step`: Can load a specific adapter checkpoint iteration
+</Note>
### Checkpoint Loading StrictnessAnd around lines 133–136:
Megatron Bridge creates checkpoints with the following directory structure:
-</Note>
+```
checkpoint_dir/🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@fern/v0.2.0/pages/training/checkpointing.mdx` around lines 81 - 135, Close
the open <Note> tag immediately before the directory tree so the "Checkpoint
Loading Strictness", "Fine-tuning Configuration", "Checkpoint Format",
"Performance Optimizations", and "Checkpoint Contents" sections are outside the
note (i.e., add the missing </Note> right after the explanatory paragraph that
precedes the directory listing), and add a matching opening code fence (```)
immediately before the directory tree so the existing closing fence correctly
encloses the tree; ensure you don't duplicate or leave an extra closing fence.
Signed-off-by: Lawrence Lane <llane@nvidia.com>
yaoyu-33
added
area:build
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary by CodeRabbit
New Features
Chores