diff --git a/megatron/core/file_arg_utils.py b/megatron/core/file_arg_utils.py new file mode 100644 index 00000000000..0e48c5777e0 --- /dev/null +++ b/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() diff --git a/megatron/core/quantization/utils.py b/megatron/core/quantization/utils.py index 5a013be27e4..f07cf19e11c 100644 --- a/megatron/core/quantization/utils.py +++ b/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 diff --git a/tests/unit_tests/core/quantization/test_load_quantization_recipe.py b/tests/unit_tests/core/quantization/test_load_quantization_recipe.py new file mode 100644 index 00000000000..22b13863916 --- /dev/null +++ b/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 diff --git a/tests/unit_tests/core/test_file_arg_utils.py b/tests/unit_tests/core/test_file_arg_utils.py new file mode 100644 index 00000000000..9a6cee1387a --- /dev/null +++ b/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"))