Spaces:
Runtime error
Runtime error
File size: 3,312 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 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
import os.path
import sys
from glob import glob
from nemo.collections.nlp.modules.common.megatron.megatron_init import (
fake_initialize_model_parallel,
)
from nemo.utils.app_state import AppState
from nemo.utils.model_utils import inject_model_parallel_rank
from omegaconf.omegaconf import OmegaConf
from trlx.data.configs import TrainConfig
from trlx.data.default_configs import default_ilql_config
from trlx.trainer.nemo_ilql_trainer import ILQLGPT, megatron_trainer
default_config = default_ilql_config()
trl_config = default_config.evolve(
train=TrainConfig(
**dict(
default_config.train.__dict__,
trainer="NeMoILQLTrainer",
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):
ilql_config = trl_config.method
megatron_cfg = OmegaConf.load(megatron_cfg_path)
megatron_cfg.trainer.num_nodes = 1
megatron_cfg.trainer.devices = 4
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)
# Manually set up the TP and PP groups
app_state = AppState()
app_state.model_parallel_size = (
megatron_cfg.model.tensor_model_parallel_size * megatron_cfg.model.pipeline_model_parallel_size
)
app_state.tensor_model_parallel_size = megatron_cfg.model.tensor_model_parallel_size
app_state.pipeline_model_parallel_size = megatron_cfg.model.pipeline_model_parallel_size
(
app_state.tensor_model_parallel_rank,
app_state.pipeline_model_parallel_rank,
app_state.model_parallel_size,
app_state.data_parallel_size,
app_state.pipeline_model_parallel_split_rank,
app_state.virtual_pipeline_model_parallel_rank,
) = fake_initialize_model_parallel(
world_size=app_state.model_parallel_size,
rank=trainer.global_rank,
tensor_model_parallel_size_=megatron_cfg.model.tensor_model_parallel_size,
pipeline_model_parallel_size_=megatron_cfg.model.pipeline_model_parallel_size,
pipeline_model_parallel_split_rank_=None,
)
checkpoint_names = find_checkpoints(checkpoint_path)
checkpoint_name = next(iter(checkpoint_names))
print(f"Loading checkpoint {checkpoint_name}, found {checkpoint_names} checkpoints")
checkpoint_path = inject_model_parallel_rank(os.path.join(checkpoint_path, checkpoint_name))
model = ILQLGPT.load_from_checkpoint(
checkpoint_path,
cfg=megatron_cfg.model,
trainer=trainer,
ilql_config=ilql_config,
)
model.sequence_parallel_(False)
model.activation_checkpointing_(False)
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])
|