Spaces:
Runtime error
Runtime error
File size: 2,160 Bytes
fa6856c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
import os.path
import sys
from glob import glob
from omegaconf.omegaconf import OmegaConf
from trlx.data.default_configs import default_ppo_config
from trlx.trainer.nemo_ppo_trainer import PPOGPT, megatron_trainer
default_config = default_ppo_config()
trl_config = default_config.evolve(
train=dict(
default_config.train.__dict__,
trainer="NeMoPPOTrainer",
trainer_kwargs=dict(
pretrained_model=None,
megatron_cfg="megatron_20b.yaml",
),
),
)
def find_checkpoints(checkpoint_dir):
checkpoints = glob(os.path.join(checkpoint_dir, "*", "*.ckpt"))
names = [os.path.basename(c) for c in checkpoints]
return set(names)
def main(megatron_cfg_path, checkpoint_path):
ppo_config = trl_config.method
megatron_cfg = OmegaConf.load(megatron_cfg_path)
megatron_cfg.trainer.num_nodes = 1
megatron_cfg.trainer.devices = (
megatron_cfg.model.tensor_model_parallel_size * megatron_cfg.model.pipeline_model_parallel_size
)
# Overriden in generate
megatron_cfg.model.global_batch_size = megatron_cfg.model.micro_batch_size
megatron_cfg.model.resume_from_checkpoint = checkpoint_path
megatron_cfg.exp_manager.create_wandb_logger = False
megatron_cfg.exp_manager.create_checkpoint_callback = False
trainer = megatron_trainer(megatron_cfg)
if trainer.world_size != megatron_cfg.trainer.devices:
raise ValueError("Inference only supports data parallel world size of 1")
# Initialize PyTorch Lightning DDP
def dummy():
return
if trainer.strategy.launcher is not None:
trainer.strategy.launcher.launch(dummy, trainer=trainer)
trainer.strategy.setup_environment()
model = PPOGPT(ppo_config=ppo_config, cfg=megatron_cfg.model, trainer=trainer, build_reference_model=False)
model.load_from_pretrained(checkpoint_path)
test = ["I don't know much about Hungarian underground"]
test = [model.tokenizer.tokenizer.bos_token + t for t in test]
print(model.generate(test, dict(max_length=40, min_length=0))["sentences"])
if __name__ == "__main__":
main(sys.argv[1], sys.argv[2])
|