Fix missing DEFAULT_INIT_MUTABLE_LIST

google · May 14, 2024 · b6add3c · b6add3c
1 parent c91d866
commit b6add3c
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/paxml/contrib/gpu/scripts_gpu/te_helper.py b/paxml/contrib/gpu/scripts_gpu/te_helper.py
@@ -1,13 +1,17 @@
 import os
 from contextlib import contextmanager
 
+from praxis import base_layer
+
 try:
   import transformer_engine.jax as te
   from transformer_engine.common import recipe
   _IS_TRANSFORMER_ENGINE_INSTALLED = True
+  DEFAULT_INIT_MUTABLE_LIST = base_layer.DEFAULT_INIT_MUTABLE_LIST + [te.fp8.FP8Helper.FP8_COLLECTION_NAME]
 
 except ModuleNotFoundError as e:
   _IS_TRANSFORMER_ENGINE_INSTALLED = False
+  DEFAULT_INIT_MUTABLE_LIST = base_layer.DEFAULT_INIT_MUTABLE_LIST
 
 
 class TransformerEngineHelperBase: