Support inline base64 payloads for te-precision-config-file

megatron/core/file_arg_utils.py

@@ -0,0 +1,17 @@
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+
+import base64
+from pathlib import Path
+
+PSEUDO_FILE_PREFIX = "base64:"
+
+
+def resolve_file_arg(value: str) -> str:
+    """Read a command line argument that is either a file path or an inline `base64:` payload.
+
+    The inline form lets a job launcher pass a config document through the command line
+    itself, without requiring a filesystem shared with the training process.
+    """
+    if value.startswith(PSEUDO_FILE_PREFIX):
+        return base64.b64decode(value[len(PSEUDO_FILE_PREFIX) :]).decode()
+    return Path(value).read_text()

megatron/core/quantization/utils.py

@@ -3,6 +3,7 @@
 import re
 from typing import Optional, Union
 
+from ..file_arg_utils import PSEUDO_FILE_PREFIX, resolve_file_arg
 from .quant_config import GlobMatcher, MatchContext, QuantizationConfig, RecipeConfig
 
 
@@ -21,7 +22,11 @@ def get_quant_config_or_none(
 
 
 def load_quantization_recipe(recipe_path: str) -> RecipeConfig:
-    """Loads a quantization recipe from a path."""
+    """Loads a quantization recipe from a path or an inline `base64:` payload."""
+    if recipe_path.startswith(PSEUDO_FILE_PREFIX):
+        import yaml
+
+        return RecipeConfig.from_config_dict(yaml.safe_load(resolve_file_arg(recipe_path)))
     recipe = RecipeConfig.from_yaml_file(recipe_path)
     return recipe
 

tests/unit_tests/core/quantization/test_load_quantization_recipe.py

@@ -0,0 +1,49 @@
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+
+import base64
+
+from megatron.core.file_arg_utils import PSEUDO_FILE_PREFIX
+from megatron.core.quantization.utils import load_quantization_recipe
+
+RECIPE_YAML = """
+configs:
+  fp8_e4m3:
+    kitchen_config_type: QLinearParams
+    quant_algo: FP8_CS
+matchers:
+  attn_qkv_fp8:
+    config: "fp8_e4m3"
+    type: "glob"
+    pattern: "*.linear_qkv"
+    enabled: true
+"""
+
+
+class TestLoadQuantizationRecipe:
+    def test_loads_a_recipe_from_a_path(self, tmp_path):
+        """The pre-existing path form must keep working for launchers that share a filesystem."""
+        path = tmp_path / "recipe.yaml"
+        path.write_text(RECIPE_YAML)
+
+        recipe = load_quantization_recipe(str(path))
+
+        assert "fp8_e4m3" in recipe.configs
+
+    def test_loads_the_same_recipe_from_an_inline_payload(self, tmp_path):
+        """The base64: wire format is what miles sends; a path and a payload must agree exactly."""
+        path = tmp_path / "recipe.yaml"
+        path.write_text(RECIPE_YAML)
+        encoded = base64.b64encode(RECIPE_YAML.encode()).decode()
+
+        from_payload = load_quantization_recipe(f"{PSEUDO_FILE_PREFIX}{encoded}")
+        from_path = load_quantization_recipe(str(path))
+
+        assert from_payload.configs == from_path.configs
+
+    def test_an_inline_recipe_keeps_its_matchers(self, tmp_path):
+        """Dropping matchers would silently quantize nothing while the recipe still looks loaded."""
+        encoded = base64.b64encode(RECIPE_YAML.encode()).decode()
+
+        recipe = load_quantization_recipe(f"{PSEUDO_FILE_PREFIX}{encoded}")
+
+        assert len(recipe.matchers) == 1

tests/unit_tests/core/test_file_arg_utils.py

@@ -0,0 +1,34 @@
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+
+import base64
+
+import pytest
+
+from megatron.core.file_arg_utils import PSEUDO_FILE_PREFIX, resolve_file_arg
+
+
+class TestResolveFileArg:
+    def test_reads_a_plain_file_path(self, tmp_path):
+        """A bare path keeps working, so existing launchers are unaffected."""
+        path = tmp_path / "recipe.yaml"
+        path.write_text("configs:\n")
+
+        assert resolve_file_arg(str(path)) == "configs:\n"
+
+    def test_decodes_an_inline_base64_payload(self):
+        """The inline form needs no filesystem shared with the job launcher."""
+        encoded = base64.b64encode(b"configs:\n").decode()
+
+        assert resolve_file_arg(f"{PSEUDO_FILE_PREFIX}{encoded}") == "configs:\n"
+
+    def test_round_trips_multiline_utf8_content(self):
+        """Recipes are multi-line and may carry non-ascii comments."""
+        text = "configs:\n  bf16:\n    # 中文注释\n"
+        encoded = base64.b64encode(text.encode()).decode()
+
+        assert resolve_file_arg(f"{PSEUDO_FILE_PREFIX}{encoded}") == text
+
+    def test_a_missing_path_still_raises(self, tmp_path):
+        """A typo must fail loudly rather than silently yield an empty recipe."""
+        with pytest.raises(FileNotFoundError):
+            resolve_file_arg(str(tmp_path / "absent.yaml"))