Rename strip_xla_flags to persist_xla_flags and add validation.

PiperOrigin-RevId: 834824081

Author: jerryxyj

export/orbax/export/constants.py

@@ -97,8 +97,13 @@ class ExportModelType(enum.Enum):
 # Mesh for the model.
 JAX_MESH = 'jax_mesh'
 
-# Whether to strip XLA flags from the model.
-STRIP_XLA_FLAGS = 'strip_xla_flags'
+# Whether to persist XLA flags in the model.
+PERSIST_XLA_FLAGS = 'persist_xla_flags'
+
+# Whether to enable bf16 optimization for the model.
+# TODO_REGEX: b/422170690: (1): Apply this flag to the pre/post processors. (2):
+# Adding filter flags once the flag is applied to the pre/post processors.
+ENABLE_BF16_OPTIMIZATION = 'enable_bf16_optimization'
 
 ################################################################################
 # Proto field names

export/orbax/export/jax_module.py

@@ -197,6 +197,16 @@ def jax2tf_kwargs_map(self) -> Mapping[str, Any]:
         tensorflow_module.TensorFlowModule, self._export_module
     ).jax2tf_kwargs_map
 
+  @property
+  def jax2obm_kwargs(self) -> Mapping[str, Any]:
+    """Returns the jax2obm_kwargs."""
+    if self._export_version == constants.ExportModelType.TF_SAVEDMODEL:
+      raise TypeError(
+          'jax2obm_kwargs is not implemented for export version'
+          ' ExportModelType.TF_SAVEDMODEL.'
+      )
+    return cast(obm_module.ObmModule, self._export_module).jax2obm_kwargs
+
   @property
   def input_polymorphic_shape_map(self) -> Mapping[str, PyTree]:
     """Returns the polymorphic shapes."""

export/orbax/export/modules/obm_module.py

@@ -73,34 +73,40 @@ def __init__(
       )
 
     # It is possible for jax2obm_kwargs to be None if the key is present.
-    if not jax2obm_kwargs:
-      jax2obm_kwargs = {}
 
+    self._jax2obm_kwargs = jax2obm_kwargs if jax2obm_kwargs else {}
+
+    enable_bf16_optimization = self.jax2obm_kwargs.get(
+        constants.ENABLE_BF16_OPTIMIZATION, False
+    )
+
+    if enable_bf16_optimization:
+      mapped_apply_fn = utils.to_bfloat16(apply_fn)
+      self._params_args_spec = utils.to_bfloat16(params)
+    else:
+      mapped_apply_fn = apply_fn
+      self._params_args_spec = params
     (
         self._apply_fn_map,
         self.input_polymorphic_shape_map,
         self.input_polymorphic_shape_symbol_values_map,
     ) = self._normalize_apply_fn_map(
-        apply_fn,
+        mapped_apply_fn,
         input_polymorphic_shape,
         input_polymorphic_shape_symbol_values,
     )
 
-    self._jax_mesh = jax2obm_kwargs.get(constants.JAX_MESH, None)
-    self._strip_xla_flags = jax2obm_kwargs.get(constants.STRIP_XLA_FLAGS, False)
+    self._jax_mesh = self.jax2obm_kwargs.get(constants.JAX_MESH, None)
 
-    self.polymorphic_constraints = self._maybe_set_polymorphic_constraints(
-        jax2obm_kwargs
-    )
+    self.polymorphic_constraints = self._maybe_set_polymorphic_constraints()
     self._native_serialization_platforms = utils.get_lowering_platforms(
-        jax2obm_kwargs
+        self.jax2obm_kwargs
     )
-    self._params_args_spec = params
 
     self._checkpoint_path: str = None
     # Set the Orbax checkpoint path if provided in the jax2obm_kwargs.
-    self._maybe_set_orbax_checkpoint_path(jax2obm_kwargs)
-    self._load_all_checkpoint_weights = jax2obm_kwargs.get(
+    self._maybe_set_orbax_checkpoint_path(self.jax2obm_kwargs)
+    self._load_all_checkpoint_weights = self.jax2obm_kwargs.get(
         constants.LOAD_ALL_CHECKPOINT_WEIGHTS, False
     )
 
@@ -203,15 +209,9 @@ def _maybe_set_orbax_checkpoint_path(self, jax2obm_kwargs):
         else constants.DEFAULT_WEIGHTS_NAME
     )
 
-  def _maybe_set_polymorphic_constraints(
-      self, jax2obm_kwargs
-  ) -> Mapping[str, Sequence[Any]]:
+  def _maybe_set_polymorphic_constraints(self) -> Mapping[str, Sequence[Any]]:
     """Sets the polymorphic constraints for the model.
 
-    Args:
-      jax2obm_kwargs: A dictionary of kwargs to pass to the jax2obm conversion
-        library.
-
     Returns:
       A mapping of function name to polymorphic constraints.
 
@@ -221,7 +221,7 @@ def _maybe_set_polymorphic_constraints(
         size of the apply_fn_map or if a key in apply_fn_map is not found in
         polymorphic_constraints.
     """
-    polymorphic_constraints = jax2obm_kwargs.get(
+    polymorphic_constraints = self.jax2obm_kwargs.get(
         constants.POLYMORPHIC_CONSTRAINTS, None
     )
     if not isinstance(polymorphic_constraints, Mapping):
@@ -300,3 +300,8 @@ def methods(self) -> Mapping[str, Callable[..., Any]]:
   def jax_methods(self) -> Mapping[str, Callable[..., Any]]:
     """Named methods in JAX context for validation."""
     raise NotImplementedError('apply_fn_map is not implemented for ObmModule.')
+
+  @property
+  def jax2obm_kwargs(self) -> Mapping[str, Any]:
+    """Returns the jax2obm_kwargs."""
+    return self._jax2obm_kwargs

export/orbax/export/modules/obm_module_test.py

@@ -357,6 +357,32 @@ def test_obm_module_multiple_apply_fns(
             jax2obm_kwargs=jax2obm_kwargs,
         )
 
+  @parameterized.named_parameters(
+      {'testcase_name': 'enable_bf16', 'enable_bf16_optimization': True},
+      {'testcase_name': 'disable_bf16', 'enable_bf16_optimization': False},
+  )
+  def test_obm_module_bfloat16_conversion(self, enable_bf16_optimization):
+    params_spec = {
+        'w': jax.ShapeDtypeStruct((2, 2), jnp.float32),
+        'b': jax.ShapeDtypeStruct((2,), jnp.float32),
+    }
+    input_spec = {constants.DEFAULT_METHOD_KEY: 'b, ...'}
+
+    module = obm_module.ObmModule(
+        params=params_spec,
+        apply_fn=_linear,
+        input_polymorphic_shape=input_spec,
+        jax2obm_kwargs={
+            constants.ENABLE_BF16_OPTIMIZATION: enable_bf16_optimization
+        },
+    )
+
+    expected_dtype = jnp.bfloat16 if enable_bf16_optimization else jnp.float32
+    with self.subTest('test_weights_w_dtype'):
+      self.assertEqual(module.model_params['w'].dtype, expected_dtype)
+    with self.subTest('test_weights_b_dtype'):
+      self.assertEqual(module.model_params['b'].dtype, expected_dtype)
+
 
 if __name__ == '__main__':
   absltest.main()

export/orbax/export/utils.py

@@ -18,6 +18,7 @@
 import dataclasses
 import functools
 import inspect
+import jax.numpy as jnp
 import os
 from typing import Any, Callable, List, Optional, Tuple, Union
 
@@ -532,3 +533,40 @@ def get_lowering_platforms(
     )
 
   return native_serialization_platforms
+
+
+def to_bfloat16(x: Any) -> Any:
+  """Helper to convert leaves of a pytree to bfloat16.
+
+  It handles `float`, `jax.ShapeDtypeStruct`, and other array-like objects with
+  a floating point `dtype`.
+
+  Args:
+    x: The input pytree to convert.
+
+  Returns:
+    The input `x` with floating point values converted to `jnp.bfloat16`.
+  """
+
+  def _to_bfloat16_leaf(x: Any) -> Any:
+    if isinstance(x, jax.ShapeDtypeStruct) and jnp.issubdtype(
+        x.dtype, jnp.floating
+    ):
+      return jax.ShapeDtypeStruct(
+          x.shape,
+          jnp.bfloat16,
+          sharding=x.sharding,
+      )
+    if isinstance(x, jax.ShapeDtypeStruct):
+      return x
+    if hasattr(x, 'dtype') and jnp.issubdtype(x.dtype, jnp.floating):
+      return x.astype(jnp.bfloat16)
+    if isinstance(x, float):
+      return jnp.bfloat16(x)
+    return x
+
+  flattened_x, treedef = jax.tree_util.tree_flatten(x)
+  flattened_y = [
+      jax.tree_util.tree_map(_to_bfloat16_leaf, y) for y in flattened_x
+  ]
+  return jax.tree_util.tree_unflatten(treedef, flattened_y)

model/orbax/experimental/model/cli/README.md

@@ -2,6 +2,7 @@
 
 A command-line tool for inspecting Orbax models.
 
+
 ## Examples
 
 To inspect the model:

model/orbax/experimental/model/core/python/compile_options_util.py

@@ -115,7 +115,7 @@ def generate_xla_compile_options(
     native_serialization_platforms: Sequence[str] | None,
     xla_flags_per_platform: Mapping[str, Sequence[str] | None],
     jax_mesh: jax.sharding.Mesh | None = None,
-    strip_xla_flags: bool = False,
+    persist_xla_flags: bool = False,
 ) -> manifest_pb2.CompileOptionsProtoMap:
   """Sets the XLA compilation options.
 
@@ -127,7 +127,7 @@ def generate_xla_compile_options(
       which will be used to override the default XLA compilation flags.
     jax_mesh: The JAX mesh used for sharding. If None, the compile options will
       be set for a default single-replica.
-    strip_xla_flags: Whether to strip XLA flags from the compile options.
+    persist_xla_flags: Whether to persist XLA flags in the compile options.
 
   Returns:
     A `CompileOptionsProtoMap` containing the XLA compilation options per
@@ -140,6 +140,9 @@ def generate_xla_compile_options(
     ValueError: If a platform is provided for XLA flags which is not provided
       in the native serialization platforms.
     ValueError: If the supplied XLA flag overrides cannot be parsed.
+    ValueError: If `xla_flags` are provided but `persist_xla_flags` is False.
+      This ensures that the XLA flags are persisted in the compile options,
+      otherwise they would be lost, leading to unexpected behavior.
   """
   tpu_platform_name = manifest_pb2.Platform.Name(
       manifest_pb2.Platform.TPU
@@ -183,16 +186,40 @@ def generate_xla_compile_options(
       xla_flags = None
       if xla_flags_per_platform:
         xla_flags = xla_flags_per_platform.get(platform, None)
+        _validate_xla_flags_setting(xla_flags, persist_xla_flags)
       compile_environment = generate_tpu_compilation_env(xla_flags)
     compile_options_map.map[platform.lower()].CopyFrom(
         generate_compilation_options(compile_environment, jax_mesh)
     )
-  if strip_xla_flags:
+  if not persist_xla_flags:
     for compile_options in compile_options_map.map.values():
       compile_options.executable_build_options.comp_envs.Clear()
   return compile_options_map
 
 
+def _validate_xla_flags_setting(
+    xla_flags: Sequence[str] | None, persist_xla_flags: bool
+) -> None:
+  """Validates the XLA flags setting.
+
+  Ensures that if `xla_flags` are provided, `persist_xla_flags` is set to True.
+  If `xla_flags` are provided but not persisted, they would be lost, leading
+  to unexpected behavior.
+
+  Args:
+    xla_flags: A sequence of XLA flags provided for overriding. Can be None.
+    persist_xla_flags: A boolean indicating whether the XLA flags should be
+      persisted in the compile options.
+
+  Raises:
+    ValueError: If `xla_flags` is not None but `persist_xla_flags` is False.
+  """
+  if xla_flags is not None and not persist_xla_flags:
+    raise ValueError(
+        'persist_xla_flags must be True if xla_flags are provided.'
+    )
+
+
 def get_field_for_flag(flag_name: str) -> descriptor.FieldDescriptor:
   """Gets the protobuf field descriptor for a given flag name."""
   if flag_name not in _XLA_FLAG_TO_FIELD_MAP:

model/orbax/experimental/model/core/python/compile_options_util_test.py

@@ -197,6 +197,7 @@ def test_generate_tpu_compilation_env_invalid_flag_format(self):
           native_serialization_platforms=None,
           xla_flags_per_platform=None,
           expected_platforms=['tpu'],
+          persist_xla_flags=False,
       ),
       dict(
           testcase_name='no_native_serialization_platforms_with_xla_flags',
@@ -205,12 +206,14 @@ def test_generate_tpu_compilation_env_invalid_flag_format(self):
               'tpu': [f'--{k}={v}' for k, v in XLA_FLAGS_DICT.items()]
           },
           expected_platforms=['tpu'],
+          persist_xla_flags=True,
       ),
       dict(
           testcase_name='with_native_serialization_platforms_no_xla_flags',
           native_serialization_platforms=['cpu', 'tpu', 'cuda'],
           xla_flags_per_platform=None,
           expected_platforms=['cpu', 'tpu', 'cuda'],
+          persist_xla_flags=False,
       ),
       dict(
           testcase_name='with_native_serialization_platforms_with_xla_flags',
@@ -219,25 +222,28 @@ def test_generate_tpu_compilation_env_invalid_flag_format(self):
               'tpu': [f'--{k}={v}' for k, v in XLA_FLAGS_DICT.items()]
           },
           expected_platforms=['cpu', 'tpu', 'cuda'],
+          persist_xla_flags=True,
       ),
   )
   def test_generate_xla_compile_options_flags_and_platforms(
       self,
       native_serialization_platforms,
       xla_flags_per_platform,
       expected_platforms,
+      persist_xla_flags,
   ):
     compile_options_map = compile_options_util.generate_xla_compile_options(
         native_serialization_platforms=native_serialization_platforms,
         xla_flags_per_platform=xla_flags_per_platform,
+        persist_xla_flags=persist_xla_flags,
     )
     self.assertLen(compile_options_map.map, len(expected_platforms))
 
     for platform in expected_platforms:
       self.assertIn(platform, compile_options_map.map)
       compile_options = compile_options_map.map[platform]
 
-      if platform != 'tpu':
+      if platform != 'tpu' or not persist_xla_flags:
         self.assertEmpty(
             compile_options.executable_build_options.comp_envs.environments
         )
@@ -307,49 +313,17 @@ def test_generate_xla_compile_options_xla_flags_platform_not_in_native_serializa
           },
       )
 
-  @parameterized.named_parameters(
-      dict(testcase_name='strip_xla_flags_true', strip_xla_flags=True),
-      dict(testcase_name='strip_xla_flags_false', strip_xla_flags=False),
-  )
-  def test_generate_xla_compile_options_strip_xla_flags(self, strip_xla_flags):
-    xla_flags_per_platform = {
-        'tpu': [f'--{k}={v}' for k, v in XLA_FLAGS_DICT.items()]
-    }
-    compile_options_map = compile_options_util.generate_xla_compile_options(
-        native_serialization_platforms=['cpu', 'tpu', 'cuda'],
-        xla_flags_per_platform=xla_flags_per_platform,
-        strip_xla_flags=strip_xla_flags,
-    )
-    self.assertLen(compile_options_map.map, 3)
-    for platform in ['cpu', 'tpu', 'cuda']:
-      self.assertIn(platform, compile_options_map.map)
-      compile_options = compile_options_map.map[platform]
-
-      if strip_xla_flags or platform != 'tpu':
-        self.assertEmpty(
-            compile_options.executable_build_options.comp_envs.environments
-        )
-      else:
-        # For TPU platform when not stripping, it should have xla flags.
-        self.assertLen(
-            compile_options.executable_build_options.comp_envs.environments, 1
-        )
-        actual_env_proto = tpu_comp_env_pb2.TpuCompilationEnvironment()
-        compile_options.executable_build_options.comp_envs.environments[
-            0
-        ].Unpack(actual_env_proto)
-
-        expected_env_overrides = EXPECTED_ENV
-        expected_env_proto = tpu_comp_env_pb2.TpuCompilationEnvironment()
-        expected_env_proto.ParseFromString(
-            tpu_comp_env.create_default_tpu_comp_env()
-        )
-        expected_env_proto.MergeFrom(expected_env_overrides)
-
-        self.assertEqual(
-            text_format.MessageToString(actual_env_proto),
-            text_format.MessageToString(expected_env_proto),
-        )
+  def test_generate_xla_compile_options_xla_flags_no_persist_raise_error(self):
+    with self.assertRaisesWithLiteralMatch(
+        ValueError, 'persist_xla_flags must be True if xla_flags are provided.'
+    ):
+      compile_options_util.generate_xla_compile_options(
+          native_serialization_platforms=['tpu'],
+          xla_flags_per_platform={
+              'tpu': [f'--{k}={v}' for k, v in XLA_FLAGS_DICT.items()]
+          },
+          persist_xla_flags=False,
+      )
 
   @parameterized.named_parameters(
       dict(