fix zch export

Author: tiankongdeguiji

tzrec/main.py

@@ -83,7 +83,7 @@
 from tzrec.utils.fx_util import symbolic_trace
 from tzrec.utils.logging_util import ProgressLogger, logger
 from tzrec.utils.plan_util import create_planner, get_default_sharders
-from tzrec.utils.state_dict_util import init_parameters, validate_state
+from tzrec.utils.state_dict_util import init_parameters
 from tzrec.version import __version__ as tzrec_version
 
 
@@ -778,7 +778,9 @@ def _script_model(
             model.load_state_dict(state_dict, strict=False)
 
         # for mc modules, we should validate and sort mch buffers
-        validate_state(model)
+        # but we do not gather state dict now, validate_state already
+        # in post_load_state_dict hook
+        # validate_state(model)
 
         batch = next(iter(dataloader))
 

tzrec/utils/state_dict_util.py

@@ -24,7 +24,6 @@ def validate_state(model: nn.Module) -> None:
             if validate_log_flag:
                 logger.info("validate states...")
                 validate_log_flag = False
-            m.validate_state()
             if isinstance(m, MCHManagedCollisionModule):
                 # fix output_segments_tensor is a meta tensor.
                 output_segments = [
@@ -36,6 +35,7 @@ def validate_state(model: nn.Module) -> None:
                     dtype=torch.int64,
                     device=m._current_iter_tensor.device,
                 )
+            m.validate_state()
 
 
 def init_parameters(module: nn.Module, device: torch.device) -> None: