pytorch · angelayi · May 16, 2025
diff --git a/docs/source/export.rst b/docs/source/export.rst
@@ -859,3 +859,5 @@ API Reference
 .. automodule:: torch.export.experimental
 .. automodule:: torch.export.passes
 .. autofunction:: torch.export.passes.move_to_device_pass
+.. automodule:: torch.export.pt2_archive
+.. automodule:: torch.export.pt2_archive.constants
diff --git a/torch/_C/_export.pyi → torch/_C/_export/__init__.pyi b/torch/_C/_export.pyi → torch/_C/_export/__init__.pyi
@@ -1,5 +1,4 @@
 # Defined in torch/csrc/export/pybind.cpp

 class CppExportedProgram: ...
 
 def deserialize_exported_program(

diff --git a/torch/_C/_export/pt2_archive_constants.pyi b/torch/_C/_export/pt2_archive_constants.pyi
@@ -0,0 +1,22 @@
+# Defined in torch/csrc/export/pt2_archive_constants.h
+
+ARCHIVE_ROOT_NAME: str = ...
+ARCHIVE_FORMAT_PATH: str = ...
+ARCHIVE_FORMAT_VALUE: str = ...
+ARCHIVE_VERSION_PATH: str = ...
+ARCHIVE_VERSION_VALUE: str = ...
+MODELS_DIR: str = ...
+MODELS_FILENAME_FORMAT: str = ...
+AOTINDUCTOR_DIR: str = ...
+MTIA_DIR: str = ...
+WEIGHTS_DIR: str = ...
+WEIGHT_FILENAME_PREFIX: str = ...
+CONSTANTS_DIR: str = ...
+TENSOR_CONSTANT_FILENAME_PREFIX: str = ...
+CUSTOM_OBJ_FILENAME_PREFIX: str = ...
+SAMPLE_INPUTS_DIR: str = ...
+SAMPLE_INPUTS_FILENAME_FORMAT: str = ...
+EXTRA_DIR: str = ...
+MODULE_INFO_PATH: str = ...
+XL_MODEL_WEIGHTS_DIR: str = ...
+XL_MODEL_WEIGHTS_PARAM_CONFIG_PATH: str = ...
diff --git a/torch/_inductor/codecache.py b/torch/_inductor/codecache.py
@@ -89,11 +89,11 @@
 from torch._utils_internal import log_cache_bypass
 from torch.compiler import config as cconfig
 from torch.compiler._cache import CacheArtifactManager, CacheArtifactType
+from torch.export.pt2_archive.constants import CUSTOM_OBJ_FILENAME_PREFIX
 from torch.fx.experimental.symbolic_shapes import has_hint, hint_int, ShapeEnv
 from torch.utils._ordered_set import OrderedSet
 
 from .output_code import CompiledFxGraph
-from .package.pt2_archive_constants import CUSTOM_OBJ_FILENAME_PREFIX
 from .remote_cache import create_cache
 from .runtime import autotune_cache
 from .runtime.autotune_cache import AutotuneCacheBundler

diff --git a/torch/_inductor/package/package.py b/torch/_inductor/package/package.py
@@ -14,14 +14,13 @@
 from torch._inductor import config
 from torch._inductor.cpp_builder import BuildOptionsBase, CppBuilder
 from torch.export._tree_utils import reorder_kwargs
-from torch.types import FileLike
-
-from .pt2_archive_constants import (
+from torch.export.pt2_archive.constants import (
     AOTINDUCTOR_DIR,
-    ARCHIVE_VERSION,
+    ARCHIVE_VERSION_VALUE,
     CONSTANTS_DIR,
     CUSTOM_OBJ_FILENAME_PREFIX,
 )
+from torch.types import FileLike
 
 
 log = logging.getLogger(__name__)
@@ -37,7 +36,7 @@ def __enter__(self) -> Self:
         self.archive_file = zipfile.ZipFile(
             self.archive_path, "w", compression=zipfile.ZIP_STORED
         )
-        self.writestr("version", str(ARCHIVE_VERSION))
+        self.writestr("version", str(ARCHIVE_VERSION_VALUE))
         self.writestr("archive_format", "pt2")
         return self
 

diff --git a/torch/_inductor/package/pt2_archive_constants.py b/torch/_inductor/package/pt2_archive_constants.py
diff --git a/torch/csrc/export/pt2_archive_constants.h b/torch/csrc/export/pt2_archive_constants.h
@@ -0,0 +1,63 @@
+#pragma once
+
+#include <array>
+#include <string_view>
+
+namespace torch::_export::archive_spec {
+
+#define FORALL_CONSTANTS(DO)                                                   \
+  DO(ARCHIVE_ROOT_NAME, "package")                                             \
+  /* Archive format */                                                         \
+  DO(ARCHIVE_FORMAT_PATH, "archive_format")                                    \
+  DO(ARCHIVE_FORMAT_VALUE, "pt2")                                              \
+  /* Archive version */                                                        \
+  DO(ARCHIVE_VERSION_PATH, "archive_version")                                  \
+  DO(ARCHIVE_VERSION_VALUE, "0") /* Sep.4.2024: This is the initial version of \
+                                    the PT2 Archive Spec */                    \
+  /*                                                                           \
+   * ######## Note on updating ARCHIVE_VERSION_VALUE ########                  \
+   * When there is a BC breaking change to the PT2 Archive Spec,               \
+   * e.g. deleting a folder, or changing the naming convention of the          \
+   * following fields it would require bumping the ARCHIVE_VERSION_VALUE       \
+   * Archive reader would need corresponding changes to support loading both   \
+   * the current and older versions of the PT2 Archive.                        \
+   */                                                                          \
+  /* Model definitions */                                                      \
+  DO(MODELS_DIR, "models/")                                                    \
+  DO(MODELS_FILENAME_FORMAT, "models/{}.json") /* {model_name} */              \
+  /* AOTInductor artifacts */                                                  \
+  DO(AOTINDUCTOR_DIR, "data/aotinductor/")                                     \
+  /* MTIA artifacts */                                                         \
+  DO(MTIA_DIR, "data/mtia")                                                    \
+  /* weights, including parameters and buffers */                              \
+  DO(WEIGHTS_DIR, "data/weights/")                                             \
+  DO(WEIGHT_FILENAME_PREFIX, "weight_")                                        \
+  /* constants, including tensor_constants, non-persistent buffers and script  \
+   * objects */                                                                \
+  DO(CONSTANTS_DIR, "data/constants/")                                         \
+  DO(TENSOR_CONSTANT_FILENAME_PREFIX, "tensor_")                               \
+  DO(CUSTOM_OBJ_FILENAME_PREFIX, "custom_obj_")                                \
+  /* example inputs */                                                         \
+  DO(SAMPLE_INPUTS_DIR, "data/sample_inputs/")                                 \
+  DO(SAMPLE_INPUTS_FILENAME_FORMAT,                                            \
+     "data/sample_inputs/{}.pt") /* {model_name} */                            \
+  /* extra folder */                                                           \
+  DO(EXTRA_DIR, "extra/")                                                      \
+  DO(MODULE_INFO_PATH, "extra/module_info.json")                               \
+  /* xl_model_weights, this folder is used for storing per-feature-weights for \
+   * remote net data in this folder is consume by Predictor, and is not        \
+   * intended to be used by Sigmoid */                                         \
+  DO(XL_MODEL_WEIGHTS_DIR, "xl_model_weights/")                                \
+  DO(XL_MODEL_WEIGHTS_PARAM_CONFIG_PATH, "xl_model_weights/model_param_config")
+
+#define DEFINE_GLOBAL(NAME, VALUE) \
+  inline constexpr std::string_view NAME = VALUE;
+FORALL_CONSTANTS(DEFINE_GLOBAL)
+#undef DEFINE_GLOBAL
+
+#define DEFINE_ENTRY(NAME, VALUE) std::pair(#NAME, VALUE),
+inline constexpr std::array kAllConstants{FORALL_CONSTANTS(DEFINE_ENTRY)};
+#undef DEFINE_ENTRY
+
+#undef FORALL_CONSTANTS
+} // namespace torch::_export::archive_spec
diff --git a/torch/csrc/export/pybind.cpp b/torch/csrc/export/pybind.cpp
@@ -1,3 +1,4 @@
+#include <torch/csrc/export/pt2_archive_constants.h>
 #include <torch/csrc/export/pybind.h>
 #include <torch/csrc/utils/generated_serialization_types.h>
 #include <torch/csrc/utils/pybind.h>
@@ -6,17 +7,23 @@ namespace torch::_export {
 
 void initExportBindings(PyObject* module) {
   auto rootModule = py::handle(module).cast<py::module>();
-  auto m = rootModule.def_submodule("_export");
+  auto exportModule = rootModule.def_submodule("_export");
+  auto pt2ArchiveModule = exportModule.def_submodule("pt2_archive_constants");
 
   // NOLINTNEXTLINE(bugprone-unused-raii)
-  py::class_<ExportedProgram>(m, "CppExportedProgram");
+  py::class_<ExportedProgram>(exportModule, "CppExportedProgram");
 
-  m.def("deserialize_exported_program", [](const std::string& serialized) {
-    return nlohmann::json::parse(serialized).get<ExportedProgram>();
-  });
+  exportModule.def(
+      "deserialize_exported_program", [](const std::string& serialized) {
+        return nlohmann::json::parse(serialized).get<ExportedProgram>();
+      });
 
-  m.def("serialize_exported_program", [](const ExportedProgram& ep) {
+  exportModule.def("serialize_exported_program", [](const ExportedProgram& ep) {
     return nlohmann::json(ep).dump();
   });
+
+  for (const auto& entry : torch::_export::archive_spec::kAllConstants) {
+    pt2ArchiveModule.attr(entry.first) = entry.second;
+  }
 }
 } // namespace torch::_export
@@ -0,0 +1,28 @@
+# Defined in torch/csrc/export/pt2_archive_constants.h
+from torch._C._export import pt2_archive_constants
+
+
+AOTINDUCTOR_DIR: str = pt2_archive_constants.AOTINDUCTOR_DIR
+ARCHIVE_FORMAT_PATH: str = pt2_archive_constants.ARCHIVE_FORMAT_PATH
+ARCHIVE_FORMAT_VALUE: str = pt2_archive_constants.ARCHIVE_FORMAT_VALUE
+ARCHIVE_ROOT_NAME: str = pt2_archive_constants.ARCHIVE_ROOT_NAME
+ARCHIVE_VERSION_PATH: str = pt2_archive_constants.ARCHIVE_VERSION_PATH
+ARCHIVE_VERSION_VALUE: str = pt2_archive_constants.ARCHIVE_VERSION_VALUE
+CONSTANTS_DIR: str = pt2_archive_constants.CONSTANTS_DIR
+CUSTOM_OBJ_FILENAME_PREFIX: str = pt2_archive_constants.CUSTOM_OBJ_FILENAME_PREFIX
+EXTRA_DIR: str = pt2_archive_constants.EXTRA_DIR
+MODELS_DIR: str = pt2_archive_constants.MODELS_DIR
+MODELS_FILENAME_FORMAT: str = pt2_archive_constants.MODELS_FILENAME_FORMAT
+MODULE_INFO_PATH: str = pt2_archive_constants.MODULE_INFO_PATH
+MTIA_DIR: str = pt2_archive_constants.MTIA_DIR
+SAMPLE_INPUTS_DIR: str = pt2_archive_constants.SAMPLE_INPUTS_DIR
+SAMPLE_INPUTS_FILENAME_FORMAT: str = pt2_archive_constants.SAMPLE_INPUTS_FILENAME_FORMAT
+TENSOR_CONSTANT_FILENAME_PREFIX: str = (
+    pt2_archive_constants.TENSOR_CONSTANT_FILENAME_PREFIX
+)
+WEIGHT_FILENAME_PREFIX: str = pt2_archive_constants.WEIGHT_FILENAME_PREFIX
+WEIGHTS_DIR: str = pt2_archive_constants.WEIGHTS_DIR
+XL_MODEL_WEIGHTS_DIR: str = pt2_archive_constants.XL_MODEL_WEIGHTS_DIR
+XL_MODEL_WEIGHTS_PARAM_CONFIG_PATH: str = (
+    pt2_archive_constants.XL_MODEL_WEIGHTS_PARAM_CONFIG_PATH
+)