| |
|
|
| """Supervised Finetuning GPT.""" |
| import functools |
| import os |
| import sys |
|
|
| import torch |
|
|
| sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../"))) |
|
|
|
|
| from examples.post_training.modelopt.finetune import SFTDataset |
| from megatron.core import mpu |
| from megatron.post_training.arguments import add_modelopt_args |
| from megatron.post_training.checkpointing import load_modelopt_checkpoint |
| from megatron.post_training.model_builder import modelopt_gpt_mamba_builder |
| from megatron.training import get_args, get_model, get_tokenizer, initialize_megatron |
| from megatron.training.utils import print_rank_0, unwrap_model |
| from model_provider import model_provider |
|
|
|
|
| def add_extract_args(parser): |
| """Add additional arguments for feature extraction.""" |
| group = parser.add_argument_group(title='Feature extraction') |
| group.add_argument("--num-samples", type=int, default=128000, help="Number of samples.") |
| group.add_argument("--output-dir", type=str, help="Path to the output directory.") |
|
|
| add_modelopt_args(parser) |
| return parser |
|
|
| def extract_feature(dataset, model, output_dir, idx_start, idx_end): |
| os.makedirs(output_dir, exist_ok=True) |
| for i in range(idx_start + mpu.get_expert_data_parallel_rank(), idx_end, mpu.get_expert_data_parallel_world_size()): |
| file_name = "{:08d}.pt".format(i - idx_start) |
| file_path = os.path.join(output_dir, file_name) |
| if not os.path.exists(file_path): |
| input_ids = dataset[i]["input_ids"][:dataset.seq_length].unsqueeze(0).to(torch.cuda.current_device()) |
| output = model(input_ids, return_eagle_inputs=True) |
| if mpu.get_tensor_model_parallel_rank() == 0 and mpu.get_expert_model_parallel_rank() == 0: |
| torch.save(output, file_path) |
| torch.distributed.barrier() |
|
|
| if __name__ == "__main__": |
| initialize_megatron( |
| extra_args_provider=add_extract_args, |
| args_defaults={ |
| 'tokenizer_type': 'HuggingFaceTokenizer', |
| 'no_load_rng': True, |
| 'no_load_optim': True, |
| }, |
| ) |
|
|
| args = get_args() |
| tokenizer = get_tokenizer() |
| model = get_model(functools.partial(model_provider, modelopt_gpt_mamba_builder), wrap_with_ddp=False) |
|
|
| load_modelopt_checkpoint(model, strict=not args.untie_embeddings_and_output_weights) |
| print_rank_0("Done loading checkpoint") |
|
|
| unwrapped_model = unwrap_model(model)[0] |
| unwrapped_model.eval() |
|
|
| kwargs = { |
| "tokenizer": tokenizer._tokenizer, |
| "seq_length": args.seq_length, |
| |
| "hf_dataset": args.finetune_hf_dataset, |
| "num_shards": mpu.get_expert_data_parallel_world_size(), |
| "shard_index": mpu.get_expert_data_parallel_rank(), |
| } |
| sft_dataset = SFTDataset(args.num_samples, None, **kwargs) |
|
|
| extract_feature(sft_dataset, unwrapped_model, os.path.join(args.output_dir, "train"), 0, int(args.num_samples * 0.98)) |
| extract_feature(sft_dataset, unwrapped_model, os.path.join(args.output_dir, "valid"), int(args.num_samples * 0.98), int(args.num_samples * 0.99)) |
| extract_feature(sft_dataset, unwrapped_model, os.path.join(args.output_dir, "test"), int(args.num_samples * 0.99), args.num_samples) |
|
|
|
|
|
|
