fix CI

verl/utils/checkpoint/megatron_checkpoint_manager.py

@@ -16,6 +16,7 @@
 import logging
 import os
 import random
+import transformer_engine
 from collections.abc import Callable
 from dataclasses import asdict
 
@@ -28,22 +29,21 @@
 from transformers import GenerationConfig
 
 # For load optimizer dist_ckpt
-import transformer_engine
 torch.serialization.add_safe_globals([torch.optim.AdamW])
 torch.serialization.add_safe_globals([transformer_engine.pytorch.optimizers.fused_adam.FusedAdam])
 
 from verl.models.weight_loader_registry import get_weight_saver
 from verl.utils.device import get_device_name, get_torch_device
 from verl.utils.fs import is_non_local, local_mkdir_safe
 from verl.utils.logger import log_with_rank
 from verl.utils.megatron.dist_checkpointing import load_dist_checkpointing, save_dist_checkpointing
 from verl.utils.megatron_utils import (
     get_dist_checkpoint_path,
     get_hf_model_checkpoint_path,
     get_transformer_config_checkpoint_path,
 )
 
 from .checkpoint_manager import BaseCheckpointManager
 
 # Setup logging
 logger = logging.getLogger(__file__)
@@ -240,7 +240,7 @@
         generate_model: bool = True,
         generate_optimizer: bool = True,
         generate_extra: bool = True,
-        is_loading: bool = False
+        is_loading: bool = False,
     ):
         # For save dist checkpointing
         state_dict = {}
@@ -308,7 +308,7 @@
             self.should_load_model and self.use_dist_checkpointing,
             self.should_load_optimizer,
             self.should_load_extra,
-            is_loading=True
+            is_loading=True,
         )
         log_with_rank(f"Generated state dict for loading: {sharded_state_dict.keys()}", rank=self.rank, logger=logger)