Don't save AutoBridge in args (#1181)

Author: zhuzilin

slime/backends/megatron_utils/actor.py

@@ -8,7 +8,6 @@
 import ray
 import torch
 import torch.distributed as dist
-from megatron.bridge import AutoBridge
 from megatron.core import mpu
 from ray.actor import ActorHandle
 from torch_memory_saver import torch_memory_saver
@@ -67,8 +66,6 @@ def init(
             if i == dist.get_rank():
                 self.hf_config = AutoConfig.from_pretrained(args.hf_checkpoint, trust_remote_code=True)
                 self.tokenizer = AutoTokenizer.from_pretrained(self.args.hf_checkpoint, trust_remote_code=True)
-                if args.megatron_to_hf_mode == "bridge":
-                    args.bridge = AutoBridge.from_hf_pretrained(args.hf_checkpoint, trust_remote_code=True)
 
             dist.barrier(group=get_gloo_group())
 

slime/backends/megatron_utils/checkpoint.py

@@ -7,6 +7,7 @@
 from megatron.training.checkpointing import load_checkpoint as _load_checkpoint_megatron
 from megatron.training.checkpointing import save_checkpoint
 from megatron.training.global_vars import get_args
+
 from slime.utils import megatron_bridge_utils
 
 logger = logging.getLogger(__name__)
@@ -48,12 +49,15 @@ def _is_megatron_checkpoint(path: str | Path) -> bool:
 
 def _load_checkpoint_hf(ddp_model, optimizer, args, load_path: str):
     assert args.megatron_to_hf_mode == "bridge", "Only bridge mode is supported for loading HF checkpoint"
+    from megatron.bridge import AutoBridge
+
     import slime_plugins.megatron_bridge  # noqa: F401
 
     logger.info(f"Load checkpoint from HuggingFace model into Megatron (path={load_path})")
 
     with megatron_bridge_utils.patch_megatron_model(ddp_model):
-        args.bridge.load_hf_weights(ddp_model)
+        bridge = AutoBridge.from_hf_pretrained(args.hf_checkpoint, trust_remote_code=True)
+        bridge.load_hf_weights(ddp_model)
 
     # Copied from Megatron-core :: load_checkpoint (with simplifications)
     if (args.fp16 or args.bf16) and optimizer is not None:

slime/backends/megatron_utils/data.py

@@ -51,7 +51,6 @@ def get_batch(
     assert "tokens" in keys
     batch = data_iterator.get_next(keys)
 
-    packed_seq_params = None
     tokens = batch["tokens"]
     # use 0 as the pad token id should be fine?
     pad_token_id = 0

slime/backends/megatron_utils/model_provider.py

@@ -54,7 +54,10 @@ def get_model_provider_func(
     role: Literal["actor", "critic"] = "actor",
 ):
     if args.megatron_to_hf_mode == "bridge":
-        provider = args.bridge.to_megatron_provider(load_weights=False)
+        from megatron.bridge import AutoBridge
+
+        bridge = AutoBridge.from_hf_pretrained(args.hf_checkpoint, trust_remote_code=True)
+        provider = bridge.to_megatron_provider(load_weights=False)
         provider.finalize()
         return provider.provide