# Generates positive movie reviews by tuning a pretrained model on IMDB dataset # with a sentiment reward function import json import os import sys from math import floor from typing import List import torch from datasets import load_dataset from transformers import AutoModelForCausalLM, AutoTokenizer import trlx from trlx.data.default_configs import ( TRLConfig, default_nemo_1_3b_config, default_ppo_config, ) def get_positive_score(scores): "Extract value associated with a positive sentiment from pipeline's output" return dict(map(lambda x: tuple(x.values()), scores))["POSITIVE"] def main(hparams={}): # Merge sweep config with default config if given default_config = TRLConfig.update(default_ppo_config().to_dict(), hparams) cfg_name = os.environ.get("NEMO_CONFIG", "1.3B") if cfg_name == "1.3B": nemo_config = default_nemo_1_3b_config() batch_size = 16 chunk_size = 128 mini_batch_size = 16 unfrozen_layers = -1 elif cfg_name == "6.7B": nemo_config = default_nemo_1_3b_config() nemo_config.name = "megatron_gpt_6.7b" nemo_config.model.num_layers = 32 nemo_config.model.hidden_size = 4096 nemo_config.model.ffn_hidden_size = 16384 nemo_config.model.num_attention_heads = 32 batch_size = 4 mini_batch_size = 4 chunk_size = 16 unfrozen_layers = -1 elif cfg_name == "13B": nemo_config = default_nemo_1_3b_config() nemo_config.name = "megatron_gpt_13b" nemo_config.model.num_layers = 40 nemo_config.model.hidden_size = 5120 nemo_config.model.ffn_hidden_size = 20480 nemo_config.model.num_attention_heads = 40 nemo_config.model.tensor_model_parallel_size = 2 batch_size = 16 mini_batch_size = 4 chunk_size = 16 unfrozen_layers = -1 elif cfg_name == "20B": nemo_config = default_nemo_1_3b_config() nemo_config.name = "megatron_gpt_20b" nemo_config.model.num_layers = 44 nemo_config.model.hidden_size = 6144 nemo_config.model.ffn_hidden_size = 24576 nemo_config.model.num_attention_heads = 64 nemo_config.model.tensor_model_parallel_size = 4 batch_size = 16 mini_batch_size = 2 chunk_size = 16 unfrozen_layers = -1 elif cfg_name == "33B": nemo_config = default_nemo_1_3b_config() nemo_config.name = "megatron_gpt_33b" nemo_config.model.num_layers = 48 nemo_config.model.hidden_size = 7168 nemo_config.model.ffn_hidden_size = 28672 nemo_config.model.num_attention_heads = 56 nemo_config.trainer.num_nodes = 4 nemo_config.trainer.devices = 8 nemo_config.model.tensor_model_parallel_size = 8 batch_size = 32 mini_batch_size = 4 chunk_size = 32 unfrozen_layers = -1 elif cfg_name == "66B": nemo_config = default_nemo_1_3b_config() nemo_config.trainer.num_nodes = 4 nemo_config.trainer.devices = 8 nemo_config.name = "megatron_gpt_66b" nemo_config.model.num_layers = 64 nemo_config.model.hidden_size = 9216 nemo_config.model.ffn_hidden_size = 36864 nemo_config.model.num_attention_heads = 72 nemo_config.model.tensor_model_parallel_size = 8 batch_size = 32 mini_batch_size = 2 chunk_size = 32 unfrozen_layers = 32 else: raise ValueError(f"Unknown NEMO_CONFIG: {cfg_name}") config = default_config.evolve( train=dict( # set automatically total_steps=None, seq_length=512, batch_size=batch_size, minibatch_size=mini_batch_size, epochs=int(1e6), eval_interval=1e6, trainer="NeMoPPOTrainer", trainer_kwargs=dict( pretrained_model=None, # f"/mnt/hdd/nemo-megatron-gpt-{cfg_name}/", megatron_cfg=nemo_config, ), checkpoint_interval=1e6, checkpoint_dir=f"nemo_{cfg_name}_ppo_ds_chat_benchmark", seed=2023, project_name="trlxnemo", tags=["nemo", "ppo", "benchmark", cfg_name], ), optimizer=dict( name="distributed_fused_adam", kwargs=dict( lr=6.001e-5, weight_decay=1e-06, eps=1.0e-8, betas=(0.9, 0.95), ), ), scheduler=dict( name="CosineAnnealing", ), model=dict(num_layers_unfrozen=unfrozen_layers), method=dict( num_rollouts=chunk_size, init_kl_coef=0.05, scale_reward="ref", vf_coef=1, gen_kwargs=dict(temperature=1.0, max_new_tokens=256, min_new_tokens=256), chunk_size=chunk_size, ppo_epochs=1, ), ) config.scheduler.kwargs = dict(warmup_steps=0, constant_steps=1e12, min_lr=6.0e-5) rank = int(os.environ["SLURM_PROCID"]) local_rank = rank % 8 reward_model = AutoModelForCausalLM.from_pretrained("facebook/opt-350m") reward_tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m") reward_model.eval() reward_model.to("cpu") def reward_fn(samples: List[str], **kwargs) -> List[float]: reward_model.to(local_rank) mbs = max(1, config.method.chunk_size // 2) for i in range(0, len(samples) // mbs): inputs = reward_tokenizer(samples[i * mbs : (i + 1) * mbs], return_tensors="pt", padding=True) inputs = inputs.to(local_rank) with torch.no_grad(): outputs = reward_model(**inputs) outputs.logits.cpu() reward_model.to("cpu") return [0.5 for _ in samples] # Take few words off of movies reviews as prompts dataset = load_dataset("Dahoas/rm-static", "train") dataset = dataset.shuffle(seed=2023) # select first 40% of the dataset dataset = dataset["train"].select(range(floor(len(dataset["train"]) * 0.4))) world_size = nemo_config.trainer.num_nodes * nemo_config.trainer.devices dp_world_size = world_size // ( nemo_config.model.tensor_model_parallel_size * nemo_config.model.pipeline_model_parallel_size ) global_batch_size = config.train.batch_size * dp_world_size config.train.total_steps = len(dataset) // global_batch_size print(f"Total steps: {config.train.total_steps=} {len(dataset)=} {global_batch_size=}") trlx.train( reward_fn=reward_fn, prompts=dataset["prompt"], eval_prompts=["I don't know much about Hungarian underground"] * 256, config=config, ) if __name__ == "__main__": hparams = {} if len(sys.argv) == 1 else json.loads(sys.argv[1]) main(hparams)