Merge branch 'main' into QAT-Walkthrough-Notebook

Author: farshadghodsian

.github/CODEOWNERS

@@ -32,24 +32,24 @@ modelopt/torch/utils @NVIDIA/modelopt-torch-utils-codeowners
 # Examples
 /docker @NVIDIA/modelopt-docker-codeowners
 /README.md @NVIDIA/modelopt-examples-codeowners
-examples @NVIDIA/modelopt-examples-codeowners
-examples/chained_optimizations @NVIDIA/modelopt-torch-nas-prune-codeowners
-examples/cnn_qat @NVIDIA/modelopt-examples-cnn_qat-codeowners
-examples/deepseek @NVIDIA/modelopt-deploy-codeowners
-examples/diffusers @NVIDIA/modelopt-examples-diffusers-codeowners
-examples/gpt-oss @NVIDIA/modelopt-examples-gpt-oss-codeowners
-examples/llm_autodeploy @NVIDIA/modelopt-deploy-codeowners
-examples/llm_distill @NVIDIA/modelopt-torch-distill-codeowners
-examples/llm_eval @NVIDIA/modelopt-examples-llm_ptq-codeowners
-examples/llm_ptq @NVIDIA/modelopt-examples-llm_ptq-codeowners
-examples/llm_qat @NVIDIA/modelopt-examples-llm_qat-codeowners
-examples/llm_sparsity @NVIDIA/modelopt-torch-sparsity-codeowners
-examples/megatron-lm @NVIDIA/modelopt-examples-megatron-codeowners
-examples/model_hub @NVIDIA/modelopt-examples-model_hub-codeowners
-examples/nemo_run @NVIDIA/modelopt-examples-megatron-codeowners
-examples/onnx_ptq @NVIDIA/modelopt-onnx-codeowners
-examples/pruning @NVIDIA/modelopt-torch-nas-prune-codeowners
-examples/speculative_decoding @NVIDIA/modelopt-torch-speculative-codeowners
-examples/vlm_eval @NVIDIA/modelopt-examples-vlm-codeowners
-examples/vlm_ptq @NVIDIA/modelopt-examples-vlm-codeowners
-examples/windows @NVIDIA/modelopt-windows-codeowners
+/examples @NVIDIA/modelopt-examples-codeowners
+/examples/chained_optimizations @NVIDIA/modelopt-torch-nas-prune-codeowners
+/examples/cnn_qat @NVIDIA/modelopt-examples-cnn_qat-codeowners
+/examples/deepseek @NVIDIA/modelopt-deploy-codeowners
+/examples/diffusers @NVIDIA/modelopt-examples-diffusers-codeowners
+/examples/gpt-oss @NVIDIA/modelopt-examples-gpt-oss-codeowners
+/examples/llm_autodeploy @NVIDIA/modelopt-deploy-codeowners
+/examples/llm_distill @NVIDIA/modelopt-torch-distill-codeowners
+/examples/llm_eval @NVIDIA/modelopt-examples-llm_ptq-codeowners
+/examples/llm_ptq @NVIDIA/modelopt-examples-llm_ptq-codeowners
+/examples/llm_qat @NVIDIA/modelopt-examples-llm_qat-codeowners
+/examples/llm_sparsity @NVIDIA/modelopt-torch-sparsity-codeowners
+/examples/megatron-lm @NVIDIA/modelopt-examples-megatron-codeowners
+/examples/model_hub @NVIDIA/modelopt-examples-model_hub-codeowners
+/examples/nemo_run @NVIDIA/modelopt-examples-megatron-codeowners
+/examples/onnx_ptq @NVIDIA/modelopt-onnx-codeowners
+/examples/pruning @NVIDIA/modelopt-torch-nas-prune-codeowners
+/examples/speculative_decoding @NVIDIA/modelopt-torch-speculative-codeowners
+/examples/vlm_eval @NVIDIA/modelopt-examples-vlm-codeowners
+/examples/vlm_ptq @NVIDIA/modelopt-examples-vlm-codeowners
+/examples/windows @NVIDIA/modelopt-windows-codeowners

.github/workflows/unit_tests.yml

@@ -4,13 +4,6 @@ name: Unit tests
 on:
   pull_request:
     branches: [main, release/*]
-    paths:
-      - ".github/workflows/unit_tests.yml"
-      - "modelopt/**"
-      - "tests/unit/**"
-      - "pyproject.toml"
-      - "setup.py"
-      - "tox.ini"
   push:
     branches: [main, release/*]
     paths:

examples/speculative_decoding/launch.sh

@@ -150,6 +150,7 @@ CMD="accelerate launch $MULTI_GPU --mixed_precision bf16 main.py \
     --logging_steps 100 \
     --tf32 True \
     --data_path $DATA \
+    --report_to tensorboard \
     $SPECULATIVE_ARGS
 "
 

tests/examples/speculative_decoding/test_eagle.py

@@ -13,12 +13,27 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import json
 
 from _test_utils.examples.run_command import run_example_command
 
 
 # fmt: off
-def test_llama_eagle(tiny_llama_path, num_gpus, tiny_daring_anteater_path, tmp_path):
+def test_llama_eagle3(tiny_llama_path, num_gpus, tiny_daring_anteater_path, tmp_path):
+    # Create an ultra-tiny EAGLE config for testing to reduce memory usage
+    tiny_eagle_config = {
+        "max_position_embeddings": 128,
+        "num_hidden_layers": 1,
+        "intermediate_size": 64,
+        "num_attention_heads": 2,
+        "num_key_value_heads": 2,
+    }
+
+    # Write the tiny config to a temporary file
+    config_file = tmp_path / "tiny_eagle_config.json"
+    with open(config_file, "w") as f:
+        json.dump(tiny_eagle_config, f)
+
     run_example_command(
         [
             "./launch.sh",
@@ -29,7 +44,9 @@ def test_llama_eagle(tiny_llama_path, num_gpus, tiny_daring_anteater_path, tmp_p
             "--do_eval", "False",
             "--num_gpu", str(num_gpus),
             "--mode", "eagle3",
+            "--eagle_config", str(config_file),
             "--output_dir", tmp_path / "eagle-tinyllama",
+            "--training_seq_len", "128", # Match max_position_embeddings
         ],
         "speculative_decoding",
     )

tests/gpu/torch/export/test_unified_export_megatron.py

@@ -110,12 +110,6 @@ def _test_unified_export_megatron(tmp_path, model_type, arch, algo, rank, size):
     ],
 )
 def test_unified_export_megatron(tmp_path, model_type, arch, algo):
-    if algo == "eagle":
-        try:
-            import megatron.core.post_training  # noqa: F401
-        except ImportError:
-            pytest.skip("megatron.core.post_training not found")
-
     # TODO: Fix TP>1 failures
     spawn_multiprocess_job(
         size=1,  # torch.cuda.device_count(),