wehub-resource-sync
128 lines
4.7 KiB
Python
128 lines
4.7 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
"""Single-model worker abstraction for Ray-based Megatron training.
|
|
|
|
MegatronModelWorker wraps one Megatron model (inference-only).
|
|
TrainableModelWorker extends it with training capabilities via _LifecycleTrainer.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import torch
|
|
from contextlib import contextmanager
|
|
from typing import Any, Dict, Optional, Sequence
|
|
|
|
from swift.utils import gc_collect, get_current_device, get_logger, to_device
|
|
|
|
logger = get_logger()
|
|
|
|
|
|
class MegatronModelWorker:
|
|
"""Wraps a single Megatron model with inference / offload interfaces.
|
|
|
|
Two creation paths:
|
|
- ``__init__(args, models)``: wraps models already created in the current
|
|
process (e.g. ref models created by MegatronRLHFTrainer.prepare_model).
|
|
- ``from_pretrained(args, model_dir)``: independently loads a new model
|
|
from disk (e.g. colocated teacher with different weights).
|
|
"""
|
|
|
|
def __init__(self, args, models, bridge=None):
|
|
self.args = args
|
|
self.models = models
|
|
self.bridge = bridge
|
|
|
|
@classmethod
|
|
def from_pretrained(cls, args, model_dir):
|
|
"""Load an inference-only model (ref / teacher) from disk."""
|
|
from transformers import AutoConfig
|
|
|
|
from swift.megatron.model import get_mcore_model
|
|
|
|
hf_config = AutoConfig.from_pretrained(model_dir, trust_remote_code=True)
|
|
models = get_mcore_model(args, hf_config)
|
|
for m in models:
|
|
if not args.use_cpu_initialization:
|
|
m.cuda(torch.cuda.current_device())
|
|
m.requires_grad_(False)
|
|
m.eval()
|
|
models[0].config.bridge.load_weights(models, model_dir)
|
|
return cls(args, models, bridge=models[0].config.bridge)
|
|
|
|
def compute_per_token_logps(self, data_iterator, temperature=1.0, enable_routing_replay=False):
|
|
from swift.megatron.trainers.utils import compute_per_token_logps_fn
|
|
return compute_per_token_logps_fn(
|
|
self.models[0],
|
|
self.args,
|
|
data_iterator,
|
|
temperature=temperature,
|
|
enable_routing_replay=enable_routing_replay)
|
|
|
|
def offload_to_cpu(self):
|
|
from swift.megatron.trainers.utils import offload_megatron_model_to_cpu
|
|
offload_megatron_model_to_cpu(self.models)
|
|
gc_collect()
|
|
|
|
def reload_to_gpu(self, load_grad=False):
|
|
from swift.megatron.trainers.utils import load_megatron_model_to_gpu
|
|
load_megatron_model_to_gpu(self.models, load_grad=load_grad)
|
|
|
|
@contextmanager
|
|
def loaded_context(self, load_grad=False):
|
|
"""Temporarily load model to GPU, offload on exit.
|
|
|
|
No-op if offloading is not configured. Use this to bracket
|
|
inference on an offloaded model (e.g. teacher forward).
|
|
"""
|
|
if not getattr(self.args, 'offload_teacher_model', False):
|
|
yield
|
|
return
|
|
self.reload_to_gpu(load_grad=load_grad)
|
|
try:
|
|
yield
|
|
finally:
|
|
self.offload_to_cpu()
|
|
|
|
|
|
class TrainableModelWorker(MegatronModelWorker):
|
|
"""Trainable model wrapping a _LifecycleTrainer with optimizer / training step.
|
|
|
|
Note: the lifecycle_trainer dependency on MegatronRLHFTrainer is a
|
|
transitional design. Future refactoring should extract the needed
|
|
capabilities (optimizer, model wrapping, ref model context) into
|
|
standalone components so the ray module no longer depends on the
|
|
non-ray trainer hierarchy.
|
|
"""
|
|
|
|
def __init__(self, args, lifecycle_trainer):
|
|
self._trainer = lifecycle_trainer
|
|
super().__init__(args, lifecycle_trainer.wrapped_models, lifecycle_trainer.bridge)
|
|
|
|
@property
|
|
def trainer(self):
|
|
return self._trainer
|
|
|
|
@property
|
|
def unwrapped_models(self):
|
|
return self._trainer.unwrapped_models
|
|
|
|
def set_forward_step(self, fn):
|
|
self._trainer.forward_step = fn
|
|
|
|
def run_train_step(self, data_iterator):
|
|
self._trainer.run_train_step(data_iterator, None)
|
|
|
|
def null_ref_context(self):
|
|
return self._trainer.null_ref_context()
|
|
|
|
def offload_to_cpu(self):
|
|
from swift.megatron.trainers.utils import offload_megatron_model_to_cpu, offload_megatron_optimizer
|
|
offload_megatron_model_to_cpu(self._trainer.wrapped_models)
|
|
if getattr(self._trainer, 'optimizer', None) and self.args.offload_optimizer:
|
|
offload_megatron_optimizer(self._trainer.optimizer)
|
|
gc_collect()
|
|
|
|
def reload_to_gpu(self, load_grad=True):
|
|
from swift.megatron.trainers.utils import load_megatron_model_to_gpu, load_megatron_optimizer
|
|
load_megatron_model_to_gpu(self._trainer.wrapped_models)
|
|
if getattr(self._trainer, 'optimizer', None) and self.args.offload_optimizer:
|
|
load_megatron_optimizer(self._trainer.optimizer)
|