From 8f45ae0afdf272c153c8456abf9d247072786c31 Mon Sep 17 00:00:00 2001 From: shh2000 <13820618441@163.com> Date: Sun, 28 Apr 2024 09:54:40 +0800 Subject: [PATCH 1/7] can run --- .../benchmarks/llama3_8B/megatron/README.md | 46 + .../llama3_8B/megatron/arguments.py | 1615 +++++++++++++++++ .../llama3_8B/megatron/pretrain_llama3.sh | 76 + .../llama3_8B/megatron/run_pretraining.py | 71 + .../llama3_8B/megatron/tokenizer.py | 567 ++++++ .../docker_image/megatron_core060/Dockerfile | 5 + .../megatron_core060_install.sh | 3 + training/nvidia/llama3_8B-megatron/README.md | 61 + .../config/config_A100x1x8.py | 7 + .../config/requirements.txt | 2 + .../config/training_adapter.sh | 1 + 11 files changed, 2454 insertions(+) create mode 100644 training/benchmarks/llama3_8B/megatron/README.md create mode 100644 training/benchmarks/llama3_8B/megatron/arguments.py create mode 100755 training/benchmarks/llama3_8B/megatron/pretrain_llama3.sh create mode 100644 training/benchmarks/llama3_8B/megatron/run_pretraining.py create mode 100644 training/benchmarks/llama3_8B/megatron/tokenizer.py create mode 100644 training/nvidia/docker_image/megatron_core060/Dockerfile create mode 100644 training/nvidia/docker_image/megatron_core060/megatron_core060_install.sh create mode 100644 training/nvidia/llama3_8B-megatron/README.md create mode 100644 training/nvidia/llama3_8B-megatron/config/config_A100x1x8.py create mode 100644 training/nvidia/llama3_8B-megatron/config/requirements.txt create mode 100644 training/nvidia/llama3_8B-megatron/config/training_adapter.sh diff --git a/training/benchmarks/llama3_8B/megatron/README.md b/training/benchmarks/llama3_8B/megatron/README.md new file mode 100644 index 000000000..1f30bc795 --- /dev/null +++ b/training/benchmarks/llama3_8B/megatron/README.md @@ -0,0 +1,46 @@ +# 模型信息 +- Introduction + +LLaMA3 is a new generation of large language models developed by Meta. The first batch of publicly released models includes two sizes: 8B and 70B parameters. We used Megatron Core-v0.6.0 and implemented the LLaMA3-8B pre-training computation task based on the Megatron framework, following the algorithm architecture and configurations of LLaMA3. + +- Technical Blog(there's no paper until Apr. 28th, 2024) + [LLAMA3](https://llama.meta.com/llama3/) + +- 模型代码来源 + +This case includes code from the LLAMA 3 COMMUNITY LICENSE AGREEMENT License open source project at:[meta-llama/Meta-Llama-3-8B · Hugging Face](https://huggingface.co/meta-llama/Meta-Llama-3-8B) + + +# 数据准备 + +### 模型配置及tokenizer准备 + +本测试样例为预训练case，需要下载tokenizer，不需要下载模型代码或模型权重。tokenizer需向llama3官方申请并下载8B版本，并在data_dir下创建llama3_8b_hf目录，按照[unsloth/llama-3-8b at main (huggingface.co)](https://huggingface.co/unsloth/llama-3-8b/tree/main)仓库示例格式存放（不需要权重部分）。 + + +### 数据集准备 + +本测试样例数据使用智源研究院wudao数据集，需向llama3官方申请tokenizer并进行预处理。测试数据集的内容不是影响AI硬件评测结果的核心因素，可以参考或使用阿里灵骏团队开放的预处理好的数据集[Pai-Megatron-Patch/examples/llama3/README.md at main · alibaba/Pai-Megatron-Patch (github.com)](https://github.com/alibaba/Pai-Megatron-Patch/blob/main/examples/llama3/README.md#%E6%95%B0%E6%8D%AE%E9%9B%86%E5%92%8C%E6%A8%A1%E5%9E%8B%E4%B8%8B%E8%BD%BD) + +将上述两个文件（.bin与.idx，不清楚文件用法可以参考Megatron-LM仓库相关介绍）放置于data_dir下。 + +# 相关开源仓库 + +### llama3（算法结构、tokenizer） + +META LLAMA 3 COMMUNITY LICENSE AGREEMENT +Meta Llama 3 Version Release Date: April 18, 2024“Agreement” means the terms and conditions for use, reproduction, distribution and modification of the Llama Materials set forth herein.“Documentation” means the specifications, manuals and documentation accompanying Meta Llama 3 distributed by Meta at https://llama.meta.com/get-started/.“Licensee” or “you” means you, or your employer or any other person or entity (if you are entering into this Agreement on such person or entity’s behalf), of the age required under applicable laws, rules or regulations to provide legal consent and that has legal authority to bind your employer or such other person or entity if you are entering in this Agreement on their behalf.“MetaLlama 3” means the foundational large language models and software and algorithms, including machine-learning model code, trained model weights, inference-enabling code, training-enabling code, fine-tuning enabling code and other elements of the foregoing distributed by Meta at https://llama.meta.com/llama-downloads.“Llama Materials” means, collectively, Meta’s proprietary Meta Llama 3 and Documentation (and any portion thereof) made available under this Agreement.“Meta” or “we” means Meta Platforms Ireland Limited (if you are located in or, if you are an entity, your principal place of business is in the EEA or Switzerland) and Meta Platforms, Inc. (if you are located outside of the EEA or Switzerland).By clicking “I Accept” below or by using or distributing any portion or element of the Llama Materials, you agree to be bound by this Agreement.1. License Rights and Redistribution.a. Grant of Rights. You are granted a non-exclusive, worldwide, non-transferable and royalty-free limited license under Meta’s intellectual property or other rights owned by Meta embodied in the Llama Materials to use, reproduce, distribute, copy, create derivative works of, and make modifications to the Llama Materials.b. Redistribution and Use.i. If you distribute or make available the Llama Materials (or any derivative works thereof), or a product or service that uses any of them, including another AI model, you shall (A) provide a copy of this Agreement with any such Llama Materials; and (B) prominently display “Built with Meta Llama 3” on a related website, user interface, blogpost, about page, or product documentation. If you use the Llama Materials to create, train, fine tune, or otherwise improve an AI model, which is distributed or made available, you shall also include “Llama 3” at the beginning of any such AI model name. +ii. If you receive Llama Materials, or any derivative works thereof, from a Licensee as part of an integrated end user product, then Section 2 of this Agreement will not apply to you. +iii. You must retain in all copies of the Llama Materials that you distribute the following attribution notice within a “Notice” text file distributed as a part of such copies: “Meta Llama 3 is licensed under the Meta Llama 3 Community License, Copyright © Meta Platforms, Inc. All Rights Reserved.”iv. Your use of the Llama Materials must comply with applicable laws and regulations (including trade compliance laws and regulations) and adhere to the Acceptable Use Policy for the Llama Materials (available at https://llama.meta.com/llama3/use-policy), which is hereby incorporated by reference into this Agreement.v. You will not use the Llama Materials or any output or results of the Llama Materials to improve any other large language model (excluding Meta Llama 3 or derivative works thereof).2. Additional Commercial Terms. If, on the Meta Llama 3 version release date, the monthly active users of the products or services made available by or for Licensee, or Licensee’s affiliates, is greater than 700 million monthly active users in the preceding calendar month, you must request a license from Meta, which Meta may grant to you in its sole discretion, and you are not authorized to exercise any of the rights under this Agreement unless or until Meta otherwise expressly grants you such rights.3. Disclaimer of Warranty. UNLESS REQUIRED BY APPLICABLE LAW, THE LLAMA MATERIALS AND ANY OUTPUT AND RESULTS THEREFROM ARE PROVIDED ON AN “AS IS” BASIS, WITHOUT WARRANTIES OF ANY KIND, AND META DISCLAIMS ALL WARRANTIES OF ANY KIND, BOTH EXPRESS AND IMPLIED, INCLUDING, WITHOUT LIMITATION, ANY WARRANTIES OF TITLE, NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. YOU ARE SOLELY RESPONSIBLE FOR DETERMINING THE APPROPRIATENESS OF USING OR REDISTRIBUTING THE LLAMA MATERIALS AND ASSUME ANY RISKS ASSOCIATED WITH YOUR USE OF THE LLAMA MATERIALS AND ANY OUTPUT AND RESULTS.4. Limitation of Liability. IN NO EVENT WILL META OR ITS AFFILIATES BE LIABLE UNDER ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, TORT, NEGLIGENCE, PRODUCTS LIABILITY, OR OTHERWISE, ARISING OUT OF THIS AGREEMENT, FOR ANY LOST PROFITS OR ANY INDIRECT, SPECIAL, CONSEQUENTIAL, INCIDENTAL, EXEMPLARY OR PUNITIVE DAMAGES, EVEN IF META OR ITS AFFILIATES HAVE BEEN ADVISED OF THE POSSIBILITY OF ANY OF THE FOREGOING.5. Intellectual Property.a. No trademark licenses are granted under this Agreement, and in connection with the Llama Materials, neither Meta nor Licensee may use any name or mark owned by or associated with the other or any of its affiliates, except as required for reasonable and customary use in describing and redistributing the Llama Materials or as set forth in this Section 5(a). Meta hereby grants you a license to use “Llama 3” (the “Mark”) solely as required to comply with the last sentence of Section 1.b.i. You will comply with Meta’s brand guidelines (currently accessible at https://about.meta.com/brand/resources/meta/company-brand/). All goodwill arising out of your use of the Mark will inure to the benefit of Meta.b. Subject to Meta’s ownership of Llama Materials and derivatives made by or for Meta, with respect to any derivative works and modifications of the Llama Materials that are made by you, as between you and Meta, you are and will be the owner of such derivative works and modifications.c. If you institute litigation or other proceedings against Meta or any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the Llama Materials or Meta Llama 3 outputs or results, or any portion of any of the foregoing, constitutes infringement of intellectual property or other rights owned or licensable by you, then any licenses granted to you under this Agreement shall terminate as of the date such litigation or claim is filed or instituted. You will indemnify and hold harmless Meta from and against any claim by any third party arising out of or related to your use or distribution of the Llama Materials.6. Term and Termination. The term of this Agreement will commence upon your acceptance of this Agreement or access to the Llama Materials and will continue in full force and effect until terminated in accordance with the terms and conditions herein. Meta may terminate this Agreement if you are in breach of any term or condition of this Agreement. Upon termination of this Agreement, you shall delete and cease use of the Llama Materials. Sections 3, 4 and 7 shall survive the termination of this Agreement.7. Governing Law and Jurisdiction. This Agreement will be governed and construed under the laws of the State of California without regard to choice of law principles, and the UN Convention on Contracts for the International Sale of Goods does not apply to this Agreement. The courts of California shall have exclusive jurisdiction of any dispute arising out of this Agreement. + +### wudao数据集（数据集内容） + +[BAAI-WuDao/WuDaoMM: WuDaoMM this is a data project (github.com)](https://github.com/BAAI-WuDao/WuDaoMM?tab=readme-ov-file) + +### 灵骏团队（megatron格式数据文件，供参考） + +[alibaba/Pai-Megatron-Patch: The official repo of Pai-Megatron-Patch for LLM & VLM large scale training developed by Alibaba Cloud. (github.com)](https://github.com/alibaba/Pai-Megatron-Patch/tree/main) + +### megatron-core（计算任务实现） + +[NVIDIA/Megatron-LM: Ongoing research training transformer models at scale (github.com)](https://github.com/NVIDIA/Megatron-LM/tree/main?tab=License-1-ov-file#readme) \ No newline at end of file diff --git a/training/benchmarks/llama3_8B/megatron/arguments.py b/training/benchmarks/llama3_8B/megatron/arguments.py new file mode 100644 index 000000000..5a062956f --- /dev/null +++ b/training/benchmarks/llama3_8B/megatron/arguments.py @@ -0,0 +1,1615 @@ +# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. + +"""Megatron arguments.""" + +import argparse +import dataclasses +import json +import os +import torch +import types + +import torch.nn.functional as F +from megatron.core.models.retro.utils import ( + get_config_path as get_retro_config_path, + get_gpt_data_dir as get_retro_data_dir, +) +from megatron.core.transformer import TransformerConfig + + +def parse_args(extra_args_provider=None, ignore_unknown_args=False): + """Parse all arguments.""" + parser = argparse.ArgumentParser(description='Megatron-LM Arguments', + allow_abbrev=False) + + # Standard arguments. + parser = _add_network_size_args(parser) + parser = _add_regularization_args(parser) + parser = _add_training_args(parser) + parser = _add_initialization_args(parser) + parser = _add_learning_rate_args(parser) + parser = _add_checkpointing_args(parser) + parser = _add_mixed_precision_args(parser) + parser = _add_distributed_args(parser) + parser = _add_validation_args(parser) + parser = _add_data_args(parser) + parser = _add_autoresume_args(parser) + parser = _add_biencoder_args(parser) + parser = _add_vision_args(parser) + parser = _add_moe_args(parser) + parser = _add_logging_args(parser) + parser = _add_inference_args(parser) + parser = _add_transformer_engine_args(parser) + parser = _add_retro_args(parser) + parser = _add_experimental_args(parser) + + # Custom arguments. + if extra_args_provider is not None: + parser = extra_args_provider(parser) + + # Parse. + if ignore_unknown_args: + args, _ = parser.parse_known_args() + else: + args = parser.parse_args() + + # Experimental yaml + if args.yaml_cfg is not None: + from .yaml_arguments import load_yaml + assert args.yaml_cfg and args.use_mcore_models, "To use yaml, mcore must be enabled" + args = load_yaml(args.yaml_cfg) + + + # Args from environment + args.rank = int(os.getenv('RANK', '0')) + args.world_size = int(os.getenv("WORLD_SIZE", '1')) + + return args + + +def load_retro_config(retro_project_dir): + '''Load Retro's config.json.''' + + # Retro config path. + retro_config_path = get_retro_config_path(retro_project_dir) + assert os.path.exists(retro_config_path), \ + "Retro project dir missing config.json." + + # Load retro config. + with open(retro_config_path) as f: + retro_config = types.SimpleNamespace(**json.load(f)) + + return retro_config + + +def load_retro_args(args): + """Load predefined args from Retro config (if applicable). + + When using Retro (or GPT for comparison purposes), data arguments are + overridden by the saved config.json within the Retro project directory. This + is to ensure that the data used for pretraining is consistent with the data + that was preprocessed using the Retro preprocessing pipeline (see + `tools/retro/preprocess_data.py`). + """ + + # Return if no project directory is specified. + if args.retro_project_dir is None: + return + + # Load retro config. + retro_config = load_retro_config(args.retro_project_dir) + + # Retro data path is relative to project dir (via hard or soft links). + data_dir = get_retro_data_dir(args.retro_project_dir) + data_path = list(retro_config.retro_gpt_data_path) + if len(data_path) % 2 == 0: + for i in range(len(data_path) - 1, -1, -2): + data_path[i] = os.path.join(data_dir, data_path[i]) + else: + assert len(data_path) == 1 + data_path[0] = os.path.join(data_dir, data_path[0]) + + # Update args. + args.data_cache_path = retro_config.retro_gpt_data_cache_path + args.data_path = data_path if args.data_path is None else args.data_path + args.eval_interval = retro_config.retro_gpt_eval_interval + args.eval_iters = retro_config.retro_gpt_eval_iters + args.global_batch_size = retro_config.retro_gpt_global_batch_size + args.max_position_embeddings = retro_config.retro_gpt_seq_length + args.merge_file = os.path.join( + args.retro_project_dir, + retro_config.retro_gpt_merge_file, + ) if retro_config.retro_gpt_merge_file is not None else None + args.seed = retro_config.retro_gpt_seed + args.seq_length = retro_config.retro_gpt_seq_length + args.tokenizer_model = os.path.join( + args.retro_project_dir, + retro_config.retro_gpt_tokenizer_model, + ) if retro_config.retro_gpt_tokenizer_model is not None else None + args.tokenizer_type = retro_config.retro_gpt_tokenizer_type + args.train_samples = retro_config.retro_gpt_train_samples + args.vocab_file = os.path.join( + args.retro_project_dir, + retro_config.retro_gpt_vocab_file, + ) if retro_config.retro_gpt_vocab_file is not None else None + + # Retro-specific args. + args.retro_block_size = retro_config.retro_block_size + args.retro_chunk_length = retro_config.retro_gpt_chunk_length + args.retro_neighbor_dirs = retro_config.retro_neighbor_dirs + args.retro_split_preprocessing = retro_config.retro_gpt_split + args.retro_bert_tokenizer_type = retro_config.retro_bert_tokenizer_type + args.retro_bert_vocab_file = retro_config.retro_bert_vocab_file + + +def validate_args(args, defaults={}): + + # Load saved args from Retro (if applicable). + load_retro_args(args) + + # Tensor model parallel size. + args.tensor_model_parallel_size = min( + args.tensor_model_parallel_size, args.world_size) + assert args.world_size % args.tensor_model_parallel_size == 0, 'world size'\ + ' ({}) is not divisible by tensor model parallel size ({})'.format( + args.world_size, args.tensor_model_parallel_size) + + # Pipeline model parallel size. + args.pipeline_model_parallel_size = min( + args.pipeline_model_parallel_size, + (args.world_size // args.tensor_model_parallel_size)) + args.transformer_pipeline_model_parallel_size = ( + args.pipeline_model_parallel_size - 1 + if args.standalone_embedding_stage else + args.pipeline_model_parallel_size + ) + + # Checks. + model_parallel_size = args.pipeline_model_parallel_size * \ + args.tensor_model_parallel_size + assert args.world_size % (model_parallel_size * args.context_parallel_size) == 0, \ + 'world size ({}) is not divisible by tensor parallel size ({}) times ' \ + 'pipeline parallel size ({}) times context parallel size ({})'.format( + args.world_size, args.tensor_model_parallel_size, + args.pipeline_model_parallel_size, args.context_parallel_size) + args.data_parallel_size = args.world_size // (model_parallel_size * args.context_parallel_size) + if args.rank == 0: + print('using world size: {}, data-parallel size: {}, ' + 'context-parallel size: {} ' + 'tensor-model-parallel size: {}, ' + 'pipeline-model-parallel size: {} '.format( + args.world_size, args.data_parallel_size, + args.context_parallel_size, + args.tensor_model_parallel_size, + args.pipeline_model_parallel_size), flush=True) + if args.pipeline_model_parallel_size > 1: + if args.pipeline_model_parallel_split_rank is not None: + assert args.pipeline_model_parallel_split_rank < \ + args.pipeline_model_parallel_size, 'split rank needs'\ + ' to be less than pipeline model parallel size ({})'.format( + args.pipeline_model_parallel_size) + + if args.tp_comm_overlap: + assert args.sequence_parallel == True, 'Tensor parallel communication/GEMM overlap can happen only when sequence parallelism is enabled' + + # Deprecated arguments + assert args.batch_size is None, '--batch-size argument is no longer ' \ + 'valid, use --micro-batch-size instead' + del args.batch_size + assert args.warmup is None, '--warmup argument is no longer valid, use ' \ + '--lr-warmup-fraction instead' + del args.warmup + assert args.model_parallel_size is None, '--model-parallel-size is no ' \ + 'longer valid, use --tensor-model-parallel-size instead' + del args.model_parallel_size + + if args.checkpoint_activations: + if args.rank == 0: + print('--checkpoint-activations is no longer valid, use --recompute-activations, ' + 'or, for more control, --recompute-granularity and --recompute-method.') + exit() + del args.checkpoint_activations + + if args.recompute_activations: + args.recompute_granularity = 'selective' + del args.recompute_activations + + # Set input defaults. + for key in defaults: + # For default to be valid, it should not be provided in the + # arguments that are passed to the program. We check this by + # ensuring the arg is set to None. + if getattr(args, key, None) is not None: + if args.rank == 0: + print('WARNING: overriding default arguments for {key}:{v} \ + with {key}:{v2}'.format(key=key, v=defaults[key], + v2=getattr(args, key)), + flush=True) + else: + setattr(args, key, defaults[key]) + + # Batch size. + assert args.micro_batch_size is not None + assert args.micro_batch_size > 0 + if args.global_batch_size is None: + args.global_batch_size = args.micro_batch_size * args.data_parallel_size + if args.rank == 0: + print('setting global batch size to {}'.format( + args.global_batch_size), flush=True) + assert args.global_batch_size > 0 + if args.num_layers_per_virtual_pipeline_stage is not None: + assert args.pipeline_model_parallel_size > 2, \ + 'pipeline-model-parallel size should be greater than 2 with ' \ + 'interleaved schedule' + assert args.num_layers % args.transformer_pipeline_model_parallel_size == 0, \ + 'number of layers should be divisible by the pipeline parallel size' + num_layers_per_pipeline_stage = args.num_layers // args.transformer_pipeline_model_parallel_size + assert num_layers_per_pipeline_stage % args.num_layers_per_virtual_pipeline_stage == 0, \ + 'number of layers per pipeline stage must be divisible number of layers per virtual pipeline stage' + args.virtual_pipeline_model_parallel_size = num_layers_per_pipeline_stage // \ + args.num_layers_per_virtual_pipeline_stage + else: + args.virtual_pipeline_model_parallel_size = None + # Overlap P2P communication is disabled if not using the interleaved schedule. + args.overlap_p2p_comm = False + if args.rank == 0: + print('WARNING: Setting args.overlap_p2p_comm to False since non-interleaved ' + 'schedule does not support overlapping p2p communication') + + if args.overlap_param_gather: + assert args.use_distributed_optimizer, \ + '--overlap-param-gather only supported with distributed optimizer' + assert args.overlap_grad_reduce, \ + '--overlap-grad-reduce should be turned on when using --overlap-param-gather' + assert args.use_mcore_models, \ + '--overlap-param-gather only supported with MCore models' + + # Parameters dtype. + args.params_dtype = torch.float + if args.fp16: + assert not args.bf16 + args.params_dtype = torch.half + # Turn off checking for NaNs in loss and grads if using dynamic loss scaling, + # where NaNs in grads / loss are signal to the loss scaler. + if not args.loss_scale: + args.check_for_nan_in_loss_and_grad = False + if args.rank == 0: + print('WARNING: Setting args.check_for_nan_in_loss_and_grad to False since ' + 'dynamic loss scaling is being used') + if args.bf16: + assert not args.fp16 + args.params_dtype = torch.bfloat16 + # bfloat16 requires gradient accumulation and all-reduce to + # be done in fp32. + if not args.accumulate_allreduce_grads_in_fp32: + args.accumulate_allreduce_grads_in_fp32 = True + if args.rank == 0: + print('accumulate and all-reduce gradients in fp32 for ' + 'bfloat16 data type.', flush=True) + + if args.rank == 0: + print('using {} for parameters ...'.format(args.params_dtype), + flush=True) + + if args.dataloader_type is None: + args.dataloader_type = 'single' + + # Consumed tokens. + args.consumed_train_samples = 0 + args.consumed_valid_samples = 0 + + # Support for variable sequence lengths across batches/microbatches. + # set it if the dataloader supports generation of variable sequence lengths + # across batches/microbatches. Due to additional communication overhead + # during pipeline parallelism, it should not be set if sequence length + # is constant during training. + args.variable_seq_lengths = False + + # Iteration-based training. + if args.train_iters: + # If we use iteration-based training, make sure the + # sample-based options are off. + assert args.train_samples is None, \ + 'expected iteration-based training' + assert args.lr_decay_samples is None, \ + 'expected iteration-based learning rate decay' + assert args.lr_warmup_samples == 0, \ + 'expected iteration-based learning rate warmup' + assert args.rampup_batch_size is None, \ + 'expected no batch-size rampup for iteration-based training' + if args.lr_warmup_fraction is not None: + assert args.lr_warmup_iters == 0, \ + 'can only specify one of lr-warmup-fraction and lr-warmup-iters' + + # Sample-based training. + if args.train_samples: + # If we use sample-based training, make sure the + # iteration-based options are off. + assert args.train_iters is None, \ + 'expected sample-based training' + assert args.lr_decay_iters is None, \ + 'expected sample-based learning rate decay' + assert args.lr_warmup_iters == 0, \ + 'expected sample-based learnig rate warmup' + if args.lr_warmup_fraction is not None: + assert args.lr_warmup_samples == 0, \ + 'can only specify one of lr-warmup-fraction ' \ + 'and lr-warmup-samples' + + if args.num_layers is not None: + assert args.encoder_num_layers is None, \ + 'cannot have both num-layers and encoder-num-layers specified' + args.encoder_num_layers = args.num_layers + else: + assert args.encoder_num_layers is not None, \ + 'either num-layers or encoder-num-layers should be specified' + args.num_layers = args.encoder_num_layers + + # Check required arguments. + required_args = ['num_layers', 'hidden_size', 'num_attention_heads', + 'max_position_embeddings'] + for req_arg in required_args: + _check_arg_is_not_none(args, req_arg) + + # Checks. + if args.ffn_hidden_size is None: + if args.swiglu: + # reduce the dimnesion for MLP since projections happens on + # two linear layers. this keeps the number of paramters in + # the same ballpark as the counterpart with 4*h size + # we keep it a multiple of 64, which means the actual tensor size + # will be a multiple of 64 / tp_size + args.ffn_hidden_size = int((4 * args.hidden_size * 2 / 3) / 64) * 64 + else: + args.ffn_hidden_size = 4 * args.hidden_size + + if args.kv_channels is None: + assert args.hidden_size % args.num_attention_heads == 0 + args.kv_channels = args.hidden_size // args.num_attention_heads + + if args.seq_length is not None: + assert args.encoder_seq_length is None + args.encoder_seq_length = args.seq_length + else: + assert args.encoder_seq_length is not None + args.seq_length = args.encoder_seq_length + + if args.seq_length is not None: + assert args.max_position_embeddings >= args.seq_length + if args.decoder_seq_length is not None: + assert args.max_position_embeddings >= args.decoder_seq_length + if args.lr is not None: + assert args.min_lr <= args.lr + if args.save is not None: + assert args.save_interval is not None + # Mixed precision checks. + if args.fp16_lm_cross_entropy: + assert args.fp16, 'lm cross entropy in fp16 only support in fp16 mode.' + if args.fp32_residual_connection: + assert args.fp16 or args.bf16, \ + 'residual connection in fp32 only supported when using fp16 or bf16.' + + if args.moe_grouped_gemm: + assert args.bf16, 'Currently GroupedGEMM for MoE only supports bf16 dtype.' + dc = torch.cuda.get_device_capability() + assert dc[0] >= 8, "Unsupported compute capability for GroupedGEMM kernels." + + if args.weight_decay_incr_style == 'constant': + assert args.start_weight_decay is None + assert args.end_weight_decay is None + args.start_weight_decay = args.weight_decay + args.end_weight_decay = args.weight_decay + else: + assert args.start_weight_decay is not None + assert args.end_weight_decay is not None + + TORCH_MAJOR = int(torch.__version__.split('.')[0]) + TORCH_MINOR = int(torch.__version__.split('.')[1]) + # Persistent fused layer norm. + if TORCH_MAJOR < 1 or (TORCH_MAJOR == 1 and TORCH_MINOR < 11): + args.no_persist_layer_norm = True + if args.rank == 0: + print('Persistent fused layer norm kernel is supported from ' + 'pytorch v1.11 (nvidia pytorch container paired with v1.11). ' + 'Defaulting to no_persist_layer_norm=True') + + # Activation recomputing. + if args.distribute_saved_activations: + assert args.tensor_model_parallel_size > 1, 'can distribute ' \ + 'recomputed activations only across tensor model ' \ + 'parallel groups' + assert args.recompute_granularity == 'full', \ + 'distributed recompute activations is only '\ + 'application to full recompute granularity' + assert args.recompute_method is not None, \ + 'for distributed recompute activations to work you '\ + 'need to use a recompute method ' + assert (TORCH_MAJOR, TORCH_MINOR) >= (1, 10), \ + 'distributed recompute activations are supported for pytorch ' \ + 'v1.10 and above (Nvidia Pytorch container >= 21.07). Current ' \ + 'pytorch version is v%s.%s.' % (TORCH_MAJOR, TORCH_MINOR) + + if args.recompute_granularity == 'selective': + assert args.recompute_method is None, \ + 'recompute method is not yet supported for ' \ + 'selective recomputing granularity' + + # disable sequence parallelism when tp=1 + # to avoid change in numerics when + # sequence_parallelism is enabled. + if args.tensor_model_parallel_size == 1: + args.sequence_parallel = False + + # disable async_tensor_model_parallel_allreduce when + # model parallel memory optimization is enabled + if args.sequence_parallel: + args.async_tensor_model_parallel_allreduce = False + + if os.environ.get('CUDA_DEVICE_MAX_CONNECTIONS') != "1": + if args.sequence_parallel: + raise RuntimeError( + "Using sequence parallelism requires setting the environment variable " + "CUDA_DEVICE_MAX_CONNECTIONS to 1") + if args.async_tensor_model_parallel_allreduce: + raise RuntimeError( + "Using async gradient all reduce requires setting the environment " + "variable CUDA_DEVICE_MAX_CONNECTIONS to 1") + + # Disable bias gelu fusion if we are disabling bias altogether + if not args.add_bias_linear: + args.bias_gelu_fusion = False + + # Retro checks. + if args.retro_add_retriever: + + # Train samples should be auto-loaded. + assert args.train_samples is not None, \ + "args.train_samples should be auto-loaded from the retro config." + + # Sequence parallelism unsupported. + assert not args.sequence_parallel, \ + "retro currently does not support sequence parallelism." + + # Pipeline parallelism unsupported. + assert args.pipeline_model_parallel_size == 1, \ + "retro currently does not support pipeline parallelism." + + if args.decoupled_lr is not None or args.decoupled_min_lr is not None: + assert args.use_mcore_models, \ + '--decoupled-lr and --decoupled-min-lr only supported by Megatron Core, please add --use-mcore-models.' + + # Legacy RoPE arguments + if args.use_rotary_position_embeddings: + args.position_embedding_type = 'rope' + if args.rotary_interleaved and args.apply_rope_fusion: + raise RuntimeError('--rotary-interleaved does not work with rope_fusion.') + if args.rotary_interleaved and not args.use_mcore_models: + raise RuntimeError('--rotary-interleaved only support Megatron Core, please add --use-mcore-models.') + + # Would just need to add 'NoPE' as a position_embedding_type to support this, but for now + # don't allow it to keep things simple + if not args.add_position_embedding and args.position_embedding_type != 'rope': + raise RuntimeError('--no-position-embedding is deprecated, use --position-embedding-type') + + # MoE Spec check + if args.num_experts is not None: + assert args.spec is None, "Model Spec must be None when using MoEs" + if args.tensor_model_parallel_size > 1: + assert args.sequence_parallel, \ + "When using MoE and tensor parallelism, sequence parallelism must be used." + + # Expert parallelism check + if args.expert_model_parallel_size > 1: + assert args.num_experts is not None, "num_experts must be non None to use expert model parallelism" + assert args.num_experts % args.expert_model_parallel_size == 0, \ + "Number of experts should be a multiple of expert model parallel_size." + assert not args.fp16, \ + "Expert parallelism is not supported with fp16 training." + + # Distributed checkpointing checks + if args.use_dist_ckpt and not args.use_mcore_models: + raise RuntimeError('--use-dist-ckpt only support Megatron Core, please add --use-mcore-models.') + + # Print arguments. + _print_args("arguments", args) + + return args + + +def _print_args(title, args): + """Print arguments.""" + if args.rank == 0: + print(f'------------------------ {title} ------------------------', + flush=True) + str_list = [] + for arg in vars(args): + dots = '.' * (48 - len(arg)) + str_list.append(' {} {} {}'.format(arg, dots, getattr(args, arg))) + for arg in sorted(str_list, key=lambda x: x.lower()): + print(arg, flush=True) + print(f'-------------------- end of {title} ---------------------', + flush=True) + + +def _check_arg_is_not_none(args, arg): + assert getattr(args, arg) is not None, '{} argument is None'.format(arg) + + +def core_transformer_config_from_args(args, config_class=None): + + # Config class. + config_class = config_class or TransformerConfig + + # Translate args to core transformer configuration + kw_args = {} + for f in dataclasses.fields(config_class): + if hasattr(args, f.name): + kw_args[f.name] = getattr(args, f.name) + kw_args['persist_layer_norm'] = not args.no_persist_layer_norm + kw_args['layernorm_zero_centered_gamma'] = args.apply_layernorm_1p + kw_args['layernorm_epsilon'] = args.norm_epsilon + kw_args['deallocate_pipeline_outputs'] = True + kw_args['pipeline_dtype'] = args.params_dtype + kw_args['batch_p2p_comm'] = not args.overlap_p2p_comm + kw_args['num_moe_experts'] = args.num_experts + kw_args['rotary_interleaved'] = args.rotary_interleaved + if args.swiglu: + kw_args['activation_func'] = F.silu + kw_args['gated_linear_unit'] = True + kw_args['bias_activation_fusion'] = args.bias_swiglu_fusion + else: + kw_args['bias_activation_fusion'] = args.bias_gelu_fusion + if args.squared_relu: + assert not args.swiglu + def squared_relu(x): + return torch.pow(F.relu(x), 2) + kw_args['activation_func'] = squared_relu + if args.init_method_xavier_uniform: + kw_args['init_method'] = torch.nn.init.xavier_uniform_ + kw_args['scaled_init_method'] = torch.nn.init.xavier_uniform_ + if args.group_query_attention: + kw_args['num_query_groups'] = args.num_query_groups + else: + kw_args['num_query_groups'] = None + + # Return config. + return config_class(**kw_args) + + +def _add_transformer_engine_args(parser): + group = parser.add_argument_group(title='Transformer-Engine') + + group.add_argument('--fp8-format', default=None, + choices=['e4m3', 'hybrid'], + help='Which fp8 format scheme to use for FP8 tensors in the forward and backward pass', + dest='fp8') + group.add_argument('--fp8-margin', type=int, default=0, + help='Scaling margin for fp8', + dest='fp8_margin') + group.add_argument('--fp8-interval', type=int, default=1, + help='Scaling update interval for fp8', + dest='fp8_interval') + group.add_argument('--fp8-amax-history-len', type=int, default=1, + help='Number of steps for which amax history is recorded per tensor', + dest='fp8_amax_history_len') + group.add_argument('--fp8-amax-compute-algo', default='most_recent', + choices=['most_recent', 'max'], + help='Algorithm for computing amax from history', + dest='fp8_amax_compute_algo') + group.add_argument('--no-fp8-wgrad', action='store_false', + help='Execute wgrad in higher precision even for FP8 runs', + dest='fp8_wgrad') + group.add_argument('--transformer-impl', default='transformer_engine', + choices=['local', 'transformer_engine'], + help='Which Transformer implementation to use.') + + return parser + +def _add_inference_args(parser): + group = parser.add_argument_group(title='inference') + + group.add_argument('--inference-batch-times-seqlen-threshold', + type=int, default=512, + help='During inference, if batch-size times ' + 'sequence-length is smaller than this threshold ' + 'then we will not use pipelining, otherwise we will.') + group.add_argument('--max-tokens-to-oom', + type=int, default=12000, + help='Maximum number of tokens during inference' + 'tokens here is # in prompt + # to generate' + 'Allows us to throw an error before OOM crashes server') + group.add_argument('--output-bert-embeddings', action='store_true', + help='Output Bert embeddings (via mean pooling) from ' + 'model, rather than its binary head output or entire ' + 'hidden batch.') + group.add_argument('--bert-embedder-type', default="megatron", + choices=["megatron", "huggingface"], + help='Select either Megatron or Huggingface as the ' + 'Bert embedder.') + + return parser + + +def _add_retro_args(parser): + group = parser.add_argument_group(title='retro') + + group.add_argument('--retro-project-dir', default=None, + help='Retro project directory, which contains the ' + 'preprocessed data for pretraining. This directory ' + 'is built during preprocessing (see ' + 'tools/retro/README.md), and contains subdirectories ' + 'for the chunk database and pretraining neighbors.') + group.add_argument('--retro-add-retriever', + action='store_true', default=False, + help='Add a retriever to the transformer, for use in ' + 'pretraining a Retro model.') + group.add_argument('--retro-cyclic-train-iters', type=int, default=None, + help='Set number of training iterations for cyclic ' + 'Retro training.') + group.add_argument('--retro-encoder-layers', type=int, default=2, + help='Number of layers to use for the retrieval ' + 'encoder.') + group.add_argument('--retro-encoder-hidden-dropout', + type=float, default=0.1, help='Hidden dropout for ' + 'retrieval encoder.') + group.add_argument('--retro-encoder-attention-dropout', + type=float, default=0.1, help='Attention dropout for ' + 'retrieval encoder.') + group.add_argument("--retro-num-neighbors", type=int, default=2, + help='Number of neighbors to retrieve during ' + 'pretraining.') + group.add_argument("--retro-num-retrieved-chunks", type=int, default=2, + help='Number of chunks to retrieve from the retrieval ' + 'database.') + group.add_argument("--retro-attention-gate", type=float, default=1, + help="Gated cross attention.") + group.add_argument("--retro-no-verify-neighbor-count", action="store_false", + dest="retro_verify_neighbor_count", + help="Skip verifying that len(GPT dataset) == len(saved " + "neighbors).") + + # Enforce argument naming convention. + for action in group._group_actions: + prefix = action.dest.split("_")[0] + assert prefix == "retro", \ + "Retro args must be prefixed with '--retro-*', for consistent " \ + "styling. Please fix '%s'." % ", ".join(action.option_strings) + + return parser + + +def _add_network_size_args(parser): + group = parser.add_argument_group(title='network size') + + group.add_argument('--num-layers', type=int, default=None, + help='Number of transformer layers.') + group.add_argument('--encoder-num-layers', type=int, default=None, + help='Number of encoder transformer layers.') + group.add_argument('--decoder-num-layers', type=int, default=None, + help='Number of decoder transformer layers.') + group.add_argument('--hidden-size', type=int, default=None, + help='Tansformer hidden size.') + group.add_argument('--ffn-hidden-size', type=int, default=None, + help='Transformer Feed-Forward Network hidden size. ' + 'This is set to 4*hidden-size if not provided') + group.add_argument('--num-attention-heads', type=int, default=None, + help='Number of transformer attention heads.') + group.add_argument('--kv-channels', type=int, default=None, + help='Projection weights dimension in multi-head ' + 'attention. This is set to ' + ' args.hidden_size // args.num_attention_heads ' + 'if not provided.') + group.add_argument('--group-query-attention', action='store_true', + help='Use group-query attention.') + group.add_argument('--num-query-groups', type=int, default=1) + + group.add_argument('--max-position-embeddings', type=int, default=None, + help='Maximum number of position embeddings to use. ' + 'This is the size of position embedding.') + group.add_argument('--position-embedding-type', type=str, default='learned_absolute', + choices=['learned_absolute', 'rope'], + help='Position embedding type.') + group.add_argument('--use-rotary-position-embeddings', action='store_true', + help='Use rotary positional embeddings or not. ' + 'Deprecated: use --position-embedding-type') + group.add_argument('--rotary-percent', type=float, default=1.0, + help='Percent of rotary dimension to use, default 100%%') + group.add_argument('--rotary-interleaved', action='store_true', + help='Use interleaved rotary embedding.') + group.add_argument('--rotary-seq-len-interpolation-factor', type=int, default=None, + help='Sequence length interpolation factor for rotary embeddings.') + group.add_argument('--no-position-embedding', + action='store_false', + help='Disable position embedding. Deprecated: use --position-embedding-type', + dest='add_position_embedding') + group.add_argument('--make-vocab-size-divisible-by', type=int, default=128, + help='Pad the vocab size to be divisible by this value.' + 'This is added for computational efficieny reasons.') + group.add_argument('--normalization', default='LayerNorm', + choices=['LayerNorm', 'RMSNorm'], + help='Which normalization technique to use.') + group.add_argument('--norm-epsilon', type=float, default=1e-5, + help='Epsilon for layer norm and RMS norm.') + group.add_argument('--apply-layernorm-1p', action='store_true', + help='Adjust LayerNorm weights such that they are centered ' + 'around zero. This improves numerical stability.') + group.add_argument('--apply-residual-connection-post-layernorm', + action='store_true', + help='If set, use original BERT residula connection ' + 'ordering.') + group.add_argument('--openai-gelu', action='store_true', + help='Use OpenAIs GeLU implementation. This option' + 'should not be used unless for backward compatibility' + 'reasons.') + group.add_argument('--squared-relu', action='store_true', + help='Use squared relu activation instead of default gelu') + group.add_argument('--swiglu', action='store_true', + help='Use gated linear units and SiLU activation instead of default gelu') + group.add_argument('--onnx-safe', type=bool, required=False, + help='Use workarounds for known problems with ' + 'Torch ONNX exporter') + group.add_argument('--bert-no-binary-head', action='store_false', + help='Disable BERT binary head.', + dest='bert_binary_head') + group.add_argument('--untie-embeddings-and-output-weights', action='store_true', + help='Untie embeddings and output weights.'), + return parser + + +def _add_logging_args(parser): + group = parser.add_argument_group(title='logging') + + group.add_argument('--log-params-norm', action='store_true', + help='If set, calculate and log parameters norm.') + group.add_argument('--log-num-zeros-in-grad', action='store_true', + help='If set, calculate and log the number of zeros in gradient.') + group.add_argument('--log-throughput', action='store_true', + help='If set, calculate and log throughput per GPU.') + group.add_argument('--log-progress', action='store_true', + help='If set, log progress (in terms of number of processed tokens and ' + 'number of floating-point operations) to progress.txt file in checkpoint ' + 'directory.') + group.add_argument('--timing-log-level', type=int, + default=0, choices=range(0,3), + help='Granularity level to measure and report timing. ' + ' 0: report only iteration time and make sure timing ' + ' does not introduce extra overhead.' + ' 1: report timing for operations that are executed ' + ' very limited times (basically once) during ' + ' each iteration (such as gradient all-reduce) ' + ' 2: report timing for operations that migh be ' + ' executed numerous times during each iteration. ' + 'Note that setting the level to 1 or 2 might ' + 'cause increase in iteration time.') + group.add_argument('--no-barrier-with-level-1-timing', action='store_false', + help='If not set, use barrier with level 1 time ' + 'measurements. Note that this is up to the user ' + 'to make sure calling barrier with their timers ' + 'will not result in hangs. This can happen if for ' + 'example the user adds a level 1 timer that is not ' + 'called by all ranks.', + dest='barrier_with_L1_time') + group.add_argument('--timing-log-option', type=str, default='minmax', + choices=['max', 'minmax', 'all'], + help='Options for logging timing:' + ' max: report the max timing across all ranks' + ' minmax: report min and max timings across all ranks' + ' all: report timings of all ranks.') + group.add_argument('--tensorboard-log-interval', type=int, default=1, + help='Report to tensorboard interval.') + group.add_argument('--tensorboard-queue-size', type=int, default=1000, + help='Size of the tensorboard queue for pending events ' + 'and summaries before one of the ‘add’ calls forces a ' + 'flush to disk.') + group.add_argument('--log-timers-to-tensorboard', action='store_true', + help='If set, write timers to tensorboard.') + group.add_argument('--log-batch-size-to-tensorboard', action='store_true', + help='If set, write batch-size to tensorboard.') + group.add_argument('--no-log-learnig-rate-to-tensorboard', + action='store_false', + help='Disable learning rate logging to tensorboard.', + dest='log_learning_rate_to_tensorboard') + group.add_argument('--no-log-loss-scale-to-tensorboard', + action='store_false', + help='Disable loss-scale logging to tensorboard.', + dest='log_loss_scale_to_tensorboard') + group.add_argument('--log-validation-ppl-to-tensorboard', + action='store_true', + help='If set, write validation perplexity to ' + 'tensorboard.') + group.add_argument('--log-memory-to-tensorboard', + action='store_true', + help='Enable memory logging to tensorboard.') + group.add_argument('--log-world-size-to-tensorboard', + action='store_true', + help='Enable world size logging to tensorboard.') + group.add_argument('--wandb-project', type=str, default='', + help='The wandb project name. Ignore wandb by default.') + group.add_argument('--wandb-exp-name', type=str, default='', + help='The wandb experiment name.') + group.add_argument('--wandb-save-dir', type=str, default='', + help='Path to save the wandb results locally.') + group.add_argument('--enable-one-logger', action='store_true', + help='If set, use one_logger to track E2E metrics' + 'Note that one_logger is an internal tool and not available externally. ' + 'For installation, please try command: `pip install ' + '--index-url=https://sc-hw-artf.nvidia.com/api/pypi/hwinf-ml-pypi/simple' + ' one_logger` or go to https://gitlab-master.nvidia.com/hwinf-dcm/onelogger ' + 'for more details') + group.add_argument('--one-logger-project', type=str, default='e2e-tracking', + help='The one-logger project name. Will ignore if ' + '--enable-one-logger is not set') + group.add_argument('--one-logger-entity', type=str, default='hwinf_dcm', + help='The one-logger username or team name. Will ignore if ' + '--enable-one-logger is not set') + group.add_argument('--one-logger-run-name', type=str, default=None, + help='The one-logger run name displayed. Will ignore if ' + '--enable-one-logger is not set') + return parser + + +def _add_regularization_args(parser): + group = parser.add_argument_group(title='regularization') + + group.add_argument('--attention-dropout', type=float, default=0.1, + help='Post attention dropout probability.') + group.add_argument('--hidden-dropout', type=float, default=0.1, + help='Dropout probability for hidden state transformer.') + group.add_argument('--weight-decay', type=float, default=0.01, + help='Weight decay coefficient for L2 regularization.') + group.add_argument('--start-weight-decay', type=float, + help='Initial weight decay coefficient for L2 regularization.') + group.add_argument('--end-weight-decay', type=float, + help='End of run weight decay coefficient for L2 regularization.') + group.add_argument('--weight-decay-incr-style', type=str, default='constant', + choices=['constant', 'linear', 'cosine'], + help='Weight decay increment function.') + group.add_argument('--clip-grad', type=float, default=1.0, + help='Gradient clipping based on global L2 norm.') + group.add_argument('--adam-beta1', type=float, default=0.9, + help='First coefficient for computing running averages ' + 'of gradient and its square') + group.add_argument('--adam-beta2', type=float, default=0.999, + help='Second coefficient for computing running averages ' + 'of gradient and its square') + group.add_argument('--adam-eps', type=float, default=1e-08, + help='Term added to the denominator to improve' + 'numerical stability') + group.add_argument('--sgd-momentum', type=float, default=0.9, + help='Momentum factor for sgd') + return parser + + +def _add_training_args(parser): + group = parser.add_argument_group(title='training') + + group.add_argument('--micro-batch-size', type=int, default=None, + help='Batch size per model instance (local batch size). ' + 'Global batch size is local batch size times data ' + 'parallel size times number of micro batches.') + group.add_argument('--batch-size', type=int, default=None, + help='Old batch size parameter, do not use. ' + 'Use --micro-batch-size instead') + group.add_argument('--global-batch-size', type=int, default=None, + help='Training batch size. If set, it should be a ' + 'multiple of micro-batch-size times data-parallel-size. ' + 'If this value is None, then ' + 'use micro-batch-size * data-parallel-size as the ' + 'global batch size. This choice will result in 1 for ' + 'number of micro-batches.') + group.add_argument('--rampup-batch-size', nargs='*', default=None, + help='Batch size ramp up with the following values:' + ' --rampup-batch-size ' + ' ' + ' ' + 'For example:' + ' --rampup-batch-size 16 8 300000 \ ' + ' --global-batch-size 1024' + 'will start with global batch size 16 and over ' + ' (1024 - 16) / 8 = 126 intervals will increase' + 'the batch size linearly to 1024. In each interval' + 'we will use approximately 300000 / 126 = 2380 samples.') + group.add_argument('--recompute-activations', action='store_true', + help='recompute activation to allow for training ' + 'with larger models, sequences, and batch sizes.') + group.add_argument('--recompute-granularity', type=str, default=None, + choices=['full', 'selective'], + help='Checkpoint activations to allow for training ' + 'with larger models, sequences, and batch sizes. ' + 'It is supported at two granularities 1) full: ' + 'whole transformer layer is recomputed, ' + '2) selective: core attention part of the transformer ' + 'layer is recomputed.') + group.add_argument('--no-check-for-nan-in-loss-and-grad', action='store_false', + help='Check for NaNs in loss and grad', + dest='check_for_nan_in_loss_and_grad') + group.add_argument('--distribute-saved-activations', + action='store_true', + help='If set, distribute recomputed activations ' + 'across model parallel group.') + group.add_argument('--recompute-method', type=str, default=None, + choices=['uniform', 'block'], + help='1) uniform: uniformly divide the total number of ' + 'Transformer layers and recompute the input activation of ' + 'each divided chunk at specified granularity, ' + '2) recompute the input activations of only a set number of ' + 'individual Transformer layers per pipeline stage and do the ' + 'rest without any recomputing at specified granularity' + 'default) do not apply activations recompute to any layers') + group.add_argument('--recompute-num-layers', type=int, default=None, + help='1) uniform: the number of Transformer layers in each ' + 'uniformly divided recompute unit, ' + '2) block: the number of individual Transformer layers ' + 'to recompute within each pipeline stage.') + group.add_argument('--no-clone-scatter-output-in-embedding', action='store_false', + help='If not set, clone the output of the scatter in embedding layer to GC original tensor.', + dest='clone_scatter_output_in_embedding') + group.add_argument('--profile', action='store_true', + help='Enable nsys profiling. When using this option, nsys ' + 'options should be specified in commandline. An example ' + 'nsys commandline is `nsys profile -s none -t nvtx,cuda ' + '-o --force-overwrite true ' + '--capture-range=cudaProfilerApi ' + '--capture-range-end=stop`.') + group.add_argument('--profile-step-start', type=int, default=10, + help='Global step to start profiling.') + group.add_argument('--profile-step-end', type=int, default=12, + help='Global step to stop profiling.') + group.add_argument('--profile-ranks', nargs='+', type=int, default=[0], + help='Global ranks to profile.') + group.add_argument('--tp-comm-overlap', action='store_true', help='Enables the ' + ' overlap of Tensor parallel communication and GEMM kernels.') + group.add_argument('--tp-comm-overlap-cfg', type=str, default=None, + help='Config file when tp_comm_overlap is enabled.') + group.add_argument('--disable-tp-comm-overlap-ag', action='store_false', + help=('Disables the All-Gather overlap with GEMM by ' + 'pipelining the GEMM and All-Gather.'), + dest='tp_comm_overlap_ag') + group.add_argument('--disable-tp-comm-overlap-rs', action='store_false', + help=('Disables the Reduce-Scatter overlap with GEMM by ' + 'pipelining the GEMM and Reduce-Scatter.'), + dest='tp_comm_overlap_rs') + group.add_argument('--disable-tp-comm-bulk-dgrad', action='store_false', + help='Disables the All-Gather overlap with bprop activation gradient GEMM.', + dest='tp_comm_bulk_dgrad') + group.add_argument('--disable-tp-comm-bulk-wgrad', action='store_false', + help='Disables the Reduce-Scatter overlap with bprop weight gradient GEMM.', + dest='tp_comm_bulk_wgrad') + group.add_argument('--use-cpu-initialization', action='store_true', + default=None, + help='If set, initialize weights on the CPU. This eliminates init differences based on tensor parallelism.') + group.add_argument('--empty-unused-memory-level', default=0, type=int, + choices=[0, 1, 2], + help='Call torch.cuda.empty_cache() each iteration ' + '(training and eval), to reduce fragmentation.' + '0=off, 1=moderate, 2=aggressive.') + + # deprecated + group.add_argument('--checkpoint-activations', action='store_true', + help='Checkpoint activation to allow for training ' + 'with larger models, sequences, and batch sizes.') + group.add_argument('--train-iters', type=int, default=None, + help='Total number of iterations to train over all ' + 'training runs. Note that either train-iters or ' + 'train-samples should be provided.') + group.add_argument('--train-samples', type=int, default=None, + help='Total number of samples to train over all ' + 'training runs. Note that either train-iters or ' + 'train-samples should be provided.') + group.add_argument('--log-interval', type=int, default=100, + help='Report loss and timing interval.') + group.add_argument('--exit-interval', type=int, default=None, + help='Exit the program after the iteration is divisible ' + 'by this value.') + group.add_argument('--exit-duration-in-mins', type=int, default=None, + help='Exit the program after this many minutes.') + group.add_argument('--exit-signal-handler', action='store_true', + help='Dynamically save the checkpoint and shutdown the ' + 'training if SIGTERM is received') + group.add_argument('--tensorboard-dir', type=str, default=None, + help='Write TensorBoard logs to this directory.') + group.add_argument('--no-masked-softmax-fusion', + action='store_false', + help='Disable fusion of query_key_value scaling, ' + 'masking, and softmax.', + dest='masked_softmax_fusion') + group.add_argument('--no-bias-gelu-fusion', action='store_false', + help='Disable bias and gelu fusion.', + dest='bias_gelu_fusion') + group.add_argument('--no-bias-swiglu-fusion', action='store_false', + help='Disable bias and swiglu fusion, the fusion is ' + 'available only when using megatron-core.', + dest='bias_swiglu_fusion') + group.add_argument('--no-bias-dropout-fusion', action='store_false', + help='Disable bias and dropout fusion.', + dest='bias_dropout_fusion') + group.add_argument('--no-rope-fusion', action='store_false', + help='Disable rope fusion, the fusion is available ' + 'only when using megatron-core.', + dest='apply_rope_fusion') + group.add_argument('--use-flash-attn', action='store_true', + help='use FlashAttention implementation of attention. ' + 'https://arxiv.org/abs/2205.14135') + group.add_argument('--disable-bias-linear', action='store_false', + help='Disable bias in the linear layers', + dest='add_bias_linear') + group.add_argument('--add-qkv-bias', action='store_true', + help='Enable bias only in the QKV linear layers', + dest='add_qkv_bias') + group.add_argument('--optimizer', type=str, default='adam', + choices=['adam', 'sgd'], + help='Optimizer function') + group.add_argument('--dataloader-type', type=str, default=None, + choices=['single', 'cyclic', 'external'], + help='Single pass vs multiple pass data loader') + group.add_argument('--no-async-tensor-model-parallel-allreduce', + action='store_false', + help='Disable asynchronous execution of ' + 'tensor-model-parallel all-reduce with weight ' + 'gradient compuation of a column-linear layer.', + dest='async_tensor_model_parallel_allreduce') + group.add_argument('--no-persist-layer-norm', action='store_true', + help='Disable using persistent fused layer norm kernel. ' + 'This kernel supports only a set of hidden sizes. Please ' + 'check persist_ln_hidden_sizes if your hidden ' + 'size is supported.') + group.add_argument('--sequence-parallel', action='store_true', + help='Enable sequence parallel optimization.') + group.add_argument('--no-gradient-accumulation-fusion', + action='store_false', + help='Disable fusing gradient accumulation to weight ' + 'gradient computation of linear layers', + dest='gradient_accumulation_fusion') + group.add_argument('--use-mcore-models', action='store_true', + help='Use the implementation from megatron core') + group.add_argument('--manual-gc', action='store_true', + help='Disable the threshold-based default garbage ' + 'collector and trigger the garbage collection manually. ' + 'Manual garbage collection helps to align the timing of ' + 'the collection across ranks which mitigates the impact ' + 'of CPU-associated jitters. When the manual gc is enabled, ' + 'garbage collection is performed only at the start and the ' + 'end of the validation routine by default.') + group.add_argument('--manual-gc-interval', type=int, default=0, + help='Training step interval to trigger manual garbage ' + 'collection. When the value is set to 0, garbage ' + 'collection is not triggered between training steps.') + group.add_argument('--no-manual-gc-eval', action='store_false', + help='When using manual garbage collection, disable ' + 'garbage collection at the start and the end of each ' + 'evaluation run.', dest='manual_gc_eval') + group.add_argument('--disable-tp-comm-split-ag', action='store_false', + help='Disables the All-Gather overlap with fprop GEMM.', + dest='tp_comm_split_ag') + group.add_argument('--disable-tp-comm-split-rs', action='store_false', + help='Disables the Reduce-Scatter overlap with fprop GEMM.', + dest='tp_comm_split_rs') + + return parser + + +def _add_initialization_args(parser): + group = parser.add_argument_group(title='initialization') + + group.add_argument('--seed', type=int, default=1234, + help='Random seed used for python, numpy, ' + 'pytorch, and cuda.') + group.add_argument('--data-parallel-random-init', action='store_true', + help='Enable random initialization of params ' + 'across data parallel ranks') + group.add_argument('--init-method-std', type=float, default=0.02, + help='Standard deviation of the zero mean normal ' + 'distribution used for weight initialization.') + group.add_argument('--init-method-xavier-uniform', action='store_true', + help='Enable Xavier uniform parameter initialization') + + return parser + + +def _add_learning_rate_args(parser): + group = parser.add_argument_group(title='learning rate') + + group.add_argument('--lr', type=float, default=None, + help='Initial learning rate. Depending on decay style ' + 'and initial warmup, the learning rate at each ' + 'iteration would be different.') + group.add_argument('--lr-decay-style', type=str, default='linear', + choices=['constant', 'linear', 'cosine', 'inverse-square-root'], + help='Learning rate decay function.') + group.add_argument('--lr-decay-iters', type=int, default=None, + help='number of iterations to decay learning rate over,' + ' If None defaults to `--train-iters`') + group.add_argument('--lr-decay-samples', type=int, default=None, + help='number of samples to decay learning rate over,' + ' If None defaults to `--train-samples`') + group.add_argument('--lr-warmup-fraction', type=float, default=None, + help='fraction of lr-warmup-(iters/samples) to use ' + 'for warmup (as a float)') + group.add_argument('--lr-warmup-iters', type=int, default=0, + help='number of iterations to linearly warmup ' + 'learning rate over.') + group.add_argument('--lr-warmup-samples', type=int, default=0, + help='number of samples to linearly warmup ' + 'learning rate over.') + group.add_argument('--lr-warmup-init', type=float, default=0.0, + help='Initial value for learning rate warmup. The ' + 'scheduler starts warmup from this value.') + group.add_argument('--warmup', type=int, default=None, + help='Old lr warmup argument, do not use. Use one of the' + '--lr-warmup-* arguments above') + group.add_argument('--min-lr', type=float, default=0.0, + help='Minimum value for learning rate. The scheduler' + 'clip values below this threshold.') + group.add_argument('--override-opt_param-scheduler', action='store_true', + help='Reset the values of the scheduler (learning rate,' + 'warmup iterations, minimum learning rate, maximum ' + 'number of iterations, and decay style from input ' + 'arguments and ignore values from checkpoints. Note' + 'that all the above values will be reset.') + group.add_argument('--use-checkpoint-opt_param-scheduler', action='store_true', + help='Use checkpoint to set the values of the scheduler ' + '(learning rate, warmup iterations, minimum learning ' + 'rate, maximum number of iterations, and decay style ' + 'from checkpoint and ignore input arguments.') + group.add_argument('--decoupled-lr', type=float, default=None, + help='Separate learning rate for the input and output layer') + group.add_argument('--decoupled-min-lr', type=float, default=None, + help='Minimum value for learning rate for the input and output layer. The scheduler' + 'clip values below this threshold') + + return parser + + +def _add_checkpointing_args(parser): + group = parser.add_argument_group(title='checkpointing') + + group.add_argument('--save', type=str, default=None, + help='Output directory to save checkpoints to.') + group.add_argument('--save-interval', type=int, default=None, + help='Number of iterations between checkpoint saves.') + group.add_argument('--no-save-optim', action='store_true', default=None, + help='Do not save current optimizer.') + group.add_argument('--no-save-rng', action='store_true', default=None, + help='Do not save current rng state.') + group.add_argument('--load', type=str, default=None, + help='Directory containing a model checkpoint.') + group.add_argument('--no-load-optim', action='store_true', default=None, + help='Do not load optimizer when loading checkpoint.') + group.add_argument('--no-load-rng', action='store_true', default=None, + help='Do not load rng state when loading checkpoint.') + group.add_argument('--finetune', action='store_true', + help='Load model for finetuning. Do not load optimizer ' + 'or rng state from checkpoint and set iteration to 0. ' + 'Assumed when loading a release checkpoint.') + group.add_argument('--pretrained-checkpoint', type=str, default=None, + help='Directory containing a pretrained model checkpoint for finetuning.') + group.add_argument('--ckpt-step', type=int, default=None, + help='Checkpoint step to load model from.') + group.add_argument('--no-initialization', action='store_false', + help='Do not perform initialization when building model, ' + 'can reduce startup time when definitely loading from a ' + 'checkpoint', + dest='perform_initialization') + group.add_argument('--use-checkpoint-args', action='store_true', + help='Override any command line arguments with arguments ' + 'from the checkpoint') + group.add_argument('--exit-on-missing-checkpoint', action='store_true', + help="If '--load' is set, but checkpoint is not found " + "(e.g., path typo), then exit instead of random " + "initialization.") + group.add_argument('--use-dist-ckpt', action='store_true', + help='Use distributed checkpoint format.') + group.add_argument('--auto-detect-ckpt-format', action='store_true', + help='Determine if the checkpoint format is in legacy or distributed format.' + ' If False, expects distributed checkpoint iff args.use_dist_ckpt.' + ' Might slow down loading a bit (double rank0 ckpt load).') + group.add_argument('--dist-ckpt-format', type=str, default='torch_dist', + choices=['zarr', 'torch_dist'], + help='Distributed checkpoint format to use.') + group.add_argument('--ckpt-fully-parallel-save', action='store_true', + help='Apply full save parallelization across DP for' + ' distributed checkpoints. Depending on ckpt format' + ' might increase number of files in the checkpoint.') + + return parser + + +def _add_mixed_precision_args(parser): + group = parser.add_argument_group(title='mixed precision') + + group.add_argument('--fp16', action='store_true', + help='Run model in fp16 mode.') + group.add_argument('--bf16', action='store_true', + help='Run model in bfloat16 mode.') + group.add_argument('--loss-scale', type=float, default=None, + help='Static loss scaling, positive power of 2 ' + 'values can improve fp16 convergence. If None, dynamic' + 'loss scaling is used.') + group.add_argument('--initial-loss-scale', type=float, default=2**32, + help='Initial loss-scale for dynamic loss scaling.') + group.add_argument('--min-loss-scale', type=float, default=1.0, + help='Minimum loss scale for dynamic loss scaling.') + group.add_argument('--loss-scale-window', type=float, default=1000, + help='Window over which to raise/lower dynamic scale.') + group.add_argument('--hysteresis', type=int, default=2, + help='hysteresis for dynamic loss scaling') + group.add_argument('--fp32-residual-connection', action='store_true', + help='Move residual connections to fp32.') + group.add_argument('--apply-query-key-layer-scaling', action='store_true', + help='Scale Q * K^T by 1 / layer-number. ' + 'Useful for fp16 training.') + group.add_argument('--attention-softmax-in-fp32', action='store_true', + help='Run attention masking and softmax in fp32. ' + 'This flag is ignored unless ' + '--no-query-key-layer-scaling is specified.') + group.add_argument('--accumulate-allreduce-grads-in-fp32', + action='store_true', + help='Gradient accumulation and all-reduce in fp32.') + group.add_argument('--fp16-lm-cross-entropy', action='store_true', + help='Move the cross entropy unreduced loss calculation' + 'for lm head to fp16.') + + return parser + + +def _add_distributed_args(parser): + group = parser.add_argument_group(title='distributed') + + group.add_argument('--tensor-model-parallel-size', type=int, default=1, + help='Degree of tensor model parallelism.') + group.add_argument('--pipeline-model-parallel-size', type=int, default=1, + help='Degree of pipeline model parallelism.') + group.add_argument('--pipeline-model-parallel-split-rank', + type=int, default=None, + help='Rank where encoder and decoder should be split.') + group.add_argument('--model-parallel-size', type=int, default=None, + help='Old model parallel argument, do not use. Use ' + '--tensor-model-parallel-size instead.') + group.add_argument('--num-layers-per-virtual-pipeline-stage', type=int, default=None, + help='Number of layers per virtual pipeline stage') + group.add_argument('--no-overlap-p2p-communication', action='store_false', + help='overlap pipeline parallel communication with forward and backward chunks', + dest='overlap_p2p_comm') + group.add_argument('--distributed-backend', default='nccl', + choices=['nccl', 'gloo'], + help='Which backend to use for distributed training.') + group.add_argument('--distributed-timeout-minutes', type=int, default=10, + help='Timeout minutes for torch.distributed.') + group.add_argument('--overlap-grad-reduce', action='store_true', + default=False, help='If set, overlap DDP grad reduce.') + group.add_argument('--no-delay-grad-reduce', action='store_false', + help='If not set, delay / synchronize grad reductions in all but first PP stage.', + dest='delay_grad_reduce') + group.add_argument('--overlap-param-gather', action='store_true', + default=False, help='If set, overlap param all-gather in distributed optimizer.') + group.add_argument('--delay-param-gather', action='store_true', + default=False, help='If set, delay / synchronize param all-gathers in all but first PP stage.') + group.add_argument('--no-scatter-gather-tensors-in-pipeline', action='store_false', + help='If not set, use scatter/gather to optimize communication of tensors in pipeline.', + dest='scatter_gather_tensors_in_pipeline') + group.add_argument('--use-ring-exchange-p2p', action='store_true', + default=False, help='If set, use custom-built ring exchange ' + 'for p2p communications. Note that this option will require ' + 'a custom built image that support ring-exchange p2p.') + group.add_argument('--local_rank', type=int, default=None, + help='local rank passed from distributed launcher.') + group.add_argument('--lazy-mpu-init', type=bool, required=False, + help='If set to True, initialize_megatron() ' + 'skips DDP initialization and returns function to ' + 'complete it instead.Also turns on ' + '--use-cpu-initialization flag. This is for ' + 'external DDP manager.' ) + group.add_argument('--standalone-embedding-stage', action='store_true', + default=False, help='If set, *input* embedding layer ' + 'is placed on its own pipeline stage, without any ' + 'transformer layers. (For T5, this flag currently only ' + 'affects the encoder embedding.)') + group.add_argument('--use-distributed-optimizer', action='store_true', + help='Use distributed optimizer.') + group.add_argument('--context-parallel-size', type=int, default=1, + help='Degree of context parallelism.') + group.add_argument('--nccl-communicator-config-path', type=str, default=None, + help='Path to the yaml file with NCCL communicator ' + 'configurations. The number of min/max thread groups and thread ' + 'group cluster size of each communicator can be configured by ' + 'setting `min_ctas`, `max_ctas`, and `cga_cluster_size`.') + return parser + + +def _add_validation_args(parser): + group = parser.add_argument_group(title='validation') + + group.add_argument('--eval-iters', type=int, default=100, + help='Number of iterations to run for evaluation' + 'validation/test for.') + group.add_argument('--eval-interval', type=int, default=1000, + help='Interval between running evaluation on ' + 'validation set.') + group.add_argument("--test-mode", action="store_true", help='Run all real-time test alongside the experiment.') + group.add_argument('--skip-train', action='store_true', + default=False, help='If set, bypass the training loop, ' + 'optionally do evaluation for validation/test, and exit.') + + return parser + + +def _add_data_args(parser): + group = parser.add_argument_group(title='data and dataloader') + + group.add_argument('--data-path', nargs='*', default=None, + help='Path to the training dataset. Accepted format:' + '1) a single data path, 2) multiple datasets in the' + 'form: dataset1-weight dataset1-path dataset2-weight ' + 'dataset2-path ... It is used with --split when a ' + 'single dataset used for all three: train, valid ' + 'and test. It is exclusive to the other ' + '--*-data-path args') + group.add_argument('--split', type=str, default='969, 30, 1', + help='Comma-separated list of proportions for training,' + ' validation, and test split. For example the split ' + '`90,5,5` will use 90%% of data for training, 5%% for ' + 'validation and 5%% for test.') + group.add_argument('--train-data-path', nargs='*', default=None, + help='Path to the training dataset. Accepted format:' + '1) a single data path, 2) multiple datasets in the' + 'form: dataset1-weight dataset1-path dataset2-weight ' + 'dataset2-path ...') + group.add_argument('--valid-data-path', nargs='*', default=None, + help='Path to the validation dataset. Accepted format:' + '1) a single data path, 2) multiple datasets in the' + 'form: dataset1-weight dataset1-path dataset2-weight ' + 'dataset2-path ...') + group.add_argument('--test-data-path', nargs='*', default=None, + help='Path to the test dataset. Accepted format:' + '1) a single data path, 2) multiple datasets in the' + 'form: dataset1-weight dataset1-path dataset2-weight ' + 'dataset2-path ...') + group.add_argument('--data-cache-path', default=None, + help='Path to a directory to hold cached index files.') + group.add_argument('--no-mmap-bin-files', action='store_false', + help='Disable mmap-ing of .bin files.', + dest='mmap_bin_files') + group.add_argument('--mock-data', action='store_true', + help='Skip data loading and validation and opt for artificial ' + 'generation of mock data when an implementation is available.') + + group.add_argument('--vocab-size', type=int, default=None, + help='Size of vocab before EOD or padding.') + group.add_argument('--vocab-file', type=str, default=None, + help='Path to the vocab file.') + group.add_argument('--merge-file', type=str, default=None, + help='Path to the BPE merge file.') + group.add_argument('--vocab-extra-ids', type=int, default=0, + help='Number of additional vocabulary tokens. ' + 'They are used for span masking in the T5 model') + group.add_argument('--seq-length', type=int, default=None, + help='Maximum sequence length to process.') + group.add_argument('--encoder-seq-length', type=int, default=None, + help='Maximum encoder sequence length to process.' + 'This should be exclusive of --seq-length') + group.add_argument('--decoder-seq-length', type=int, default=None, + help="Maximum decoder sequence length to process.") + group.add_argument('--retriever-seq-length', type=int, default=256, + help='Maximum sequence length for the biencoder model ' + 'for retriever') + group.add_argument('--sample-rate', type=float, default=1.0, + help='sample rate for training data. Supposed to be 0 ' + ' < sample_rate < 1') + group.add_argument('--mask-prob', type=float, default=0.15, + help='Probability of replacing a token with mask.') + group.add_argument('--short-seq-prob', type=float, default=0.1, + help='Probability of producing a short sequence.') + group.add_argument('--num-workers', type=int, default=2, + help="Dataloader number of workers.") + group.add_argument('--tokenizer-type', type=str, + default=None, + choices=['BertWordPieceLowerCase', + 'BertWordPieceCase', + 'GPT2BPETokenizer', + 'SentencePieceTokenizer', + 'GPTSentencePieceTokenizer', + 'Llama2Tokenizer', + 'Llama3Tokenizer', + 'NullTokenizer'], + help='What type of tokenizer to use.') + group.add_argument('--tokenizer-model', type=str, default=None, + help='Sentencepiece tokenizer model.') + group.add_argument('--reset-position-ids', action='store_true', + help='Reset posistion ids after end-of-document token.') + group.add_argument('--reset-attention-mask', action='store_true', + help='Reset self attention maske after ' + 'end-of-document token.') + group.add_argument('--eod-mask-loss', action='store_true', + help='Mask loss for the end of document tokens.') + group.add_argument('--no-create-attention-mask-in-dataloader', action='store_false', + help='If set, do not create attention_masks in dataloader.', + dest='create_attention_mask_in_dataloader') + + return parser + + +def _add_autoresume_args(parser): + group = parser.add_argument_group(title='autoresume') + + group.add_argument('--adlr-autoresume', action='store_true', + help='Enable autoresume on adlr cluster.') + group.add_argument('--adlr-autoresume-interval', type=int, default=1000, + help='Intervals over which check for autoresume' + 'termination signal') + + return parser + + +def _add_biencoder_args(parser): + group = parser.add_argument_group(title='biencoder') + + # network size + group.add_argument('--ict-head-size', type=int, default=None, + help='Size of block embeddings to be used in ICT and ' + 'REALM (paper default: 128)') + group.add_argument('--biencoder-projection-dim', type=int, default=0, + help='Size of projection head used in biencoder (paper' + ' default: 128)') + group.add_argument('--biencoder-shared-query-context-model', action='store_true', + help='Whether to share the parameters of the query ' + 'and context models or not') + + # checkpointing + group.add_argument('--ict-load', type=str, default=None, + help='Directory containing an ICTBertModel checkpoint') + group.add_argument('--bert-load', type=str, default=None, + help='Directory containing an BertModel checkpoint ' + '(needed to start ICT and REALM)') + + # data + group.add_argument('--titles-data-path', type=str, default=None, + help='Path to titles dataset used for ICT') + group.add_argument('--query-in-block-prob', type=float, default=0.1, + help='Probability of keeping query in block for ' + 'ICT dataset') + group.add_argument('--use-one-sent-docs', action='store_true', + help='Whether to use one sentence documents in ICT') + group.add_argument('--evidence-data-path', type=str, default=None, + help='Path to Wikipedia Evidence frm DPR paper') + + # training + group.add_argument('--retriever-report-topk-accuracies', nargs='+', type=int, + default=[], help="Which top-k accuracies to report " + "(e.g. '1 5 20')") + group.add_argument('--retriever-score-scaling', action='store_true', + help='Whether to scale retriever scores by inverse ' + 'square root of hidden size') + + # faiss index + group.add_argument('--block-data-path', type=str, default=None, + help='Where to save/load BlockData to/from') + group.add_argument('--embedding-path', type=str, default=None, + help='Where to save/load Open-Retrieval Embedding' + ' data to/from') + + # indexer + group.add_argument('--indexer-batch-size', type=int, default=128, + help='How large of batches to use when doing indexing ' + 'jobs') + group.add_argument('--indexer-log-interval', type=int, default=1000, + help='After how many batches should the indexer ' + 'report progress') + return parser + + +def _add_vision_args(parser): + group = parser.add_argument_group(title="vision") + + # general vision arguements + group.add_argument('--num-classes', type=int, default=1000, + help='num of classes in vision classificaiton task') + group.add_argument('--img-h', type=int, default=224, + help='Image height for vision classification task') + group.add_argument('--img-w', type=int, default=224, + help='Image height for vision classification task') + group.add_argument('--num-channels', type=int, default=3, + help='Number of channels in input image data') + group.add_argument('--patch-dim', type=int, default=16, + help='patch dimension') + group.add_argument('--classes-fraction', type=float, default=1.0, + help='training with fraction of classes.') + group.add_argument('--data-per-class-fraction', type=float, default=1.0, + help='training with fraction of data per class.') + group.add_argument('--no-data-sharding', action='store_false', + help='Disable data sharding.', + dest='data_sharding') + group.add_argument('--head-lr-mult', type=float, default=1.0, + help='learning rate multiplier for head during finetuning') + + # pretraining type and backbone selection` + group.add_argument('--vision-pretraining', action='store_true', + help='flag to indicate vision pretraining') + group.add_argument('--vision-pretraining-type', type=str, default='classify', + choices=['classify', 'inpaint', 'dino'], + help='pretraining objectives') + group.add_argument('--vision-backbone-type', type=str, default='vit', + choices=['vit', 'mit', 'swin'], + help='backbone types types') + group.add_argument('--swin-backbone-type', type=str, default='tiny', + choices=['tiny', 'base', 'h3'], + help='pretraining objectives') + # inpainting arguments + group.add_argument('--mask-type', type=str, default='random', + choices=['random', 'row'], + help='mask types') + group.add_argument('--mask-factor', type=float, default=1.0, + help='mask size scaling parameter') + + # dino arguments + group.add_argument('--iter-per-epoch', type=int, default=1250, + help='iterations per epoch') + group.add_argument('--dino-local-img-size', type=int, default=96, + help='Image size for vision classification task') + group.add_argument('--dino-local-crops-number', type=int, default=10, + help='Number of local crops') + group.add_argument('--dino-head-hidden-size', type=int, default=2048, + help='Hidden dimension size in dino head') + group.add_argument('--dino-bottleneck-size', type=int, default=256, + help='Bottle neck dimension in dino head ') + group.add_argument('--dino-freeze-last-layer', type=float, default=1, + help='Freezing last layer weights') + group.add_argument('--dino-norm-last-layer', action='store_true', + help='Disable Norm in last layer.') + group.add_argument('--dino-warmup-teacher-temp', type=float, default=0.04, + help='warump teacher temperature') + group.add_argument('--dino-teacher-temp', type=float, default=0.07, + help='teacher temperature') + group.add_argument('--dino-warmup-teacher-temp-epochs', type=int, default=30, + help='warmup teacher temperaure epochs') + + # regularization arguments + group.add_argument('--qk-layernorm', action='store_true', + help='Whether to layer normalize the q and k attention embeddings.') + + return parser + +def _add_moe_args(parser): + group = parser.add_argument_group(title="moe") + group.add_argument('--expert-model-parallel-size', type=int, default=1, + help='Degree of expert model parallelism.') + group.add_argument('--num-experts', type=int, default=None, + help='Number of Experts in MoE (None means no MoE)') + group.add_argument('--moe-router-load-balancing-type', type=str, + choices=['aux_loss', 'sinkhorn', "none"], + default='aux_loss', + help='Determines the load balancing strategy for the router. "aux_loss" corresponds to the load balancing loss used in GShard and SwitchTransformer, "sinkhorn" corresponds to the balancing algorithm used in S-BASE, and "none" implies no load balancing. The default is "aux_loss".') + group.add_argument('--moe-router-topk', type=int, default=2, + help='Number of experts to route to for each token. The default is 2.') + group.add_argument('--moe-grouped-gemm', action='store_true', + help='When there are multiple experts per rank, compress multiple local (potentially small) gemms in a single kernel launch to improve the utilization and performance by leveraging the Grouped GEMM feature introduced since CUTLASS 2.8 (https://github.com/fanshiqing/grouped_gemm).') + group.add_argument('--moe-aux-loss-coeff', type=float, default=0.0, + help='Scaling coefficient for the aux loss: a starting value of 1e-2 is recommended.') + group.add_argument('--moe-z-loss-coeff', type=float, default=None, + help='Scaling coefficient for the z-loss: a starting value of 1e-3 is recommended.') + group.add_argument('--moe-input-jitter-eps', type=float, default=None, + help='Add noise to the input tensor by applying jitter with a specified epsilon value.') + group.add_argument('--moe-token-dropping', action='store_true', + help='This feature involves selectively dropping and padding tokens for each expert to achieve a specified capacity, similar to GShard, Switch-Transformer, and DeepSpeed-MoE. Note: Currently unsupported.') + group.add_argument('--moe-token-dispatcher-type', type=str, + choices=['allgather', 'alltoall'], + default='allgather', + help='.') + group.add_argument('--moe-per-layer-logging', action='store_true', + help='Enable per-layer logging for MoE, currently supports auxiliary loss and z loss.') + + return parser + +def _add_experimental_args(parser): + group = parser.add_argument_group(title='experimental') + + group.add_argument('--spec', type=str, default=None, nargs='*', + help='Specify the pair ' + 'that returns a spec to customize a model, transformer ' + 'block, or transformer layer, depending on the use case.' + 'To use local spec specify local as the argument.' + 'For more details, see the model class, ' + '`transformer_block.py`, or `transformer_layer.py`') + group.add_argument('--yaml-cfg', type=str, default=None, + help = 'Config file to add additional arguments') + + return parser diff --git a/training/benchmarks/llama3_8B/megatron/pretrain_llama3.sh b/training/benchmarks/llama3_8B/megatron/pretrain_llama3.sh new file mode 100755 index 000000000..ee1eca7ff --- /dev/null +++ b/training/benchmarks/llama3_8B/megatron/pretrain_llama3.sh @@ -0,0 +1,76 @@ +#!/bin/bash + +# Runs the "LLaMA3-8B" parameter model +export PYTHONPATH=$PYTHONPATH:/workspace/Megatron-LM +export CUDA_DEVICE_MAX_CONNECTIONS=1 + +ls /workspace +# algorithm args + +MODEL_ARGS=" \ + --num-layers 32 \ + --hidden-size 4096 \ + --ffn-hidden-size 14336 \ + --num-attention-heads 32 \ + --group-query-attention \ + --num-query-groups 8 \ + --seq-length 8192 \ + --max-position-embeddings 8192 \ + --swiglu \ + --normalization RMSNorm \ + --norm-epsilon 1.0e-5 \ + --global-batch-size 512 \ + --attention-softmax-in-fp32" + +OPT_ARGS=" \ + --lr 1.0e-5 \ + --train-iters 300 \ + --lr-decay-iters 300 \ + --lr-decay-style cosine \ + --min-lr 1.0e-6 \ + --weight-decay 0.1 \ + --clip-grad 1.0 \ + --adam-beta1 0.9 \ + --adam-beta2 0.95 \ + --init-method-std 0.008" + +ALGO_ARGS="$MODEL_ARGS $OPT_ARGS" + +# data args + +DATA_ARGS=" \ + --data-path /data/llama3_8b_pretrain/wudao_llama3bpe_content_document \ + --tokenizer-type Llama3Tokenizer \ + --split 100,0,0 +" + +# training args + +TRAINING_ARGS=" \ + --micro-batch-size 1 \ + --tensor-model-parallel-size 1 \ + --pipeline-model-parallel-size 2 \ + --sequence-parallel \ + --bf16 +" + +# vendor args + +VENDOR_ARGS=" \ + --transformer-impl transformer_engine \ + --use-distributed-optimizer \ + --use-mcore-models \ + --use-flash-attn +" + +OUTPUT_ARGS=" --log-interval 1" + +run_cmd="torchrun --nproc_per_node 8 --nnodes 1 --node_rank 0 --master_addr 10.1.2.59 --master_port 29501 /workspace/Megatron-LM/pretrain_gpt.py \ + $ALGO_ARGS \ + $DATA_ARGS \ + $TRAINING_ARGS \ + $VENDOR_ARGS \ + $OUTPUT_ARGS" + +echo ${run_cmd} +eval ${run_cmd} diff --git a/training/benchmarks/llama3_8B/megatron/run_pretraining.py b/training/benchmarks/llama3_8B/megatron/run_pretraining.py new file mode 100644 index 000000000..e076054fe --- /dev/null +++ b/training/benchmarks/llama3_8B/megatron/run_pretraining.py @@ -0,0 +1,71 @@ +import subprocess +from argparse import ArgumentParser +import os +import sys +from importlib import import_module + + +def parse_args(): + '''we parse ddp related args, check system config args, and running env + args such as --data_dir_xxx. Then pass all useful args to the real + training script. + ''' + parser = ArgumentParser(description="flagscale main python") + parser.add_argument("--nproc_per_node", type=int, required=True) + parser.add_argument("--nnodes", type=int, required=True) + parser.add_argument("--node_rank", type=int, required=True) + parser.add_argument("--master_addr", type=str, required=True) + parser.add_argument("--master_port", type=int, required=True) + parser.add_argument("--vendor", type=str, required=True) + parser.add_argument("--data_dir", type=str, required=True) + parser.add_argument("--log_dir", type=str, required=True) + parser.add_argument("--flagperf_config_file", type=str, required=True) + args, unknown_args = parser.parse_known_args() + args.unknown_args = unknown_args + return args + + +if __name__ == "__main__": + args = parse_args() + print(args) + + sys.path.append(os.path.dirname(args.flagperf_config_file)) + config_file = os.path.basename(args.flagperf_config_file).split('.')[0] + config_dir_path = os.path.dirname(args.flagperf_config_file) + + module = import_module(config_file) + + seqlength = getattr(module, 'seqlength') + localbs = getattr(module, 'localbs') + train_steps = getattr(module, 'train_steps') + theoryflops = getattr(module, 'theoryflops') + epochs = getattr(module, 'epochs') + tensor_parallel = getattr(module, 'tensor_parallel') + pipeline_parallel = getattr(module, 'pipeline_parallel') + + + task_log_file = os.path.join(args.log_dir, "megatron.log.txt") + + exec_cmd = "cp arguments.py /workspace/Megatron-LM/megatron/training/;cp tokenizer.py /workspace/Megatron-LM/megatron/training/tokenizer/;bash pretrain_llama3.sh" + exec_cmd = exec_cmd + " " + os.path.join(config_dir_path, "training_adapter.sh") + + with open(task_log_file, "w") as f: + p = subprocess.Popen(exec_cmd, + shell=True, + stdout=f, + stderr=subprocess.STDOUT) + p.wait() + + time_per_step = -1.0 + with open(task_log_file) as f: + for line in f.readlines(): + if "elapsed time per iteration (ms): " in line: + info = line.split("|")[2] + steptime = info.split(":")[1] + time_per_step = float(steptime) / 1000 + + whole_tps = 512* seqlength / time_per_step + chip_tps = whole_tps / (args.nproc_per_node * args.nnodes) + print("System tokens per second: ", whole_tps) + print("Tokens/p/s: ", chip_tps) + print("MFU: ", chip_tps * 80000000000.0 * 6 / theoryflops) diff --git a/training/benchmarks/llama3_8B/megatron/tokenizer.py b/training/benchmarks/llama3_8B/megatron/tokenizer.py new file mode 100644 index 000000000..b824f84ea --- /dev/null +++ b/training/benchmarks/llama3_8B/megatron/tokenizer.py @@ -0,0 +1,567 @@ +# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +"""Megatron tokenizers.""" + +from abc import ABC +from abc import abstractmethod + +from megatron.core.datasets.megatron_tokenizer import MegatronTokenizer + +from .bert_tokenization import FullTokenizer as FullBertTokenizer +from .gpt2_tokenization import GPT2Tokenizer + + +def build_tokenizer(args): + """Initialize tokenizer.""" + if args.rank == 0: + print('> building {} tokenizer ...'.format(args.tokenizer_type), + flush=True) + + # Select and instantiate the tokenizer. + if args.tokenizer_type == 'BertWordPieceLowerCase': + assert args.vocab_file is not None + tokenizer = _BertWordPieceTokenizer(vocab_file=args.vocab_file, + lower_case=True, + vocab_extra_ids=args.vocab_extra_ids) + elif args.tokenizer_type == 'BertWordPieceCase': + assert args.vocab_file is not None + tokenizer = _BertWordPieceTokenizer(vocab_file=args.vocab_file, + lower_case=False, + vocab_extra_ids=args.vocab_extra_ids) + elif args.tokenizer_type == 'GPT2BPETokenizer': + assert args.vocab_file is not None + assert args.merge_file is not None + tokenizer = _GPT2BPETokenizer(args.vocab_file, args.merge_file) + elif args.tokenizer_type == 'SentencePieceTokenizer': + assert args.tokenizer_model is not None + tokenizer = _SentencePieceTokenizer(args.tokenizer_model, vocab_extra_ids=args.vocab_extra_ids) + elif args.tokenizer_type == 'GPTSentencePieceTokenizer': + assert args.tokenizer_model is not None + tokenizer = _GPTSentencePieceTokenizer(args.tokenizer_model) + elif args.tokenizer_type == 'Llama2Tokenizer': + assert args.tokenizer_model is not None + tokenizer = _Llama2Tokenizer(args.tokenizer_model) + elif args.tokenizer_type == 'Llama3Tokenizer': + tokenizer = _Llama3Tokenizer() + elif args.tokenizer_type == 'NullTokenizer': + assert args.vocab_size is not None + tokenizer = _NullTokenizer(args.vocab_size) + else: + raise NotImplementedError('{} tokenizer is not ' + 'implemented.'.format(args.tokenizer_type)) + + # Add vocab size (if not already set from a checkpoint). + if getattr(args, "padded_vocab_size", None) is None: + args.padded_vocab_size = _vocab_size_with_padding(tokenizer.vocab_size, + args) + + return tokenizer + + +def _vocab_size_with_padding(orig_vocab_size, args): + """Pad vocab size so it is divisible by model parallel size and + still having GPU friendly size.""" + + after = orig_vocab_size + multiple = args.make_vocab_size_divisible_by * \ + args.tensor_model_parallel_size + while (after % multiple) != 0: + after += 1 + if args.rank == 0: + print(' > padded vocab (size: {}) with {} dummy tokens ' + '(new size: {})'.format( + orig_vocab_size, after - orig_vocab_size, after), flush=True) + return after + + +class _BertWordPieceTokenizer(MegatronTokenizer): + """Original BERT wordpiece tokenizer.""" + + def __init__(self, vocab_file, lower_case=True, vocab_extra_ids=0): + super().__init__(vocab_file, lower_case=lower_case, vocab_extra_ids=vocab_extra_ids) + self.tokenizer = FullBertTokenizer(vocab_file, do_lower_case=lower_case) + self.cls_id = self.tokenizer.vocab['[CLS]'] + self.sep_id = self.tokenizer.vocab['[SEP]'] + self.pad_id = self.tokenizer.vocab['[PAD]'] + self.mask_id = self.tokenizer.vocab['[MASK]'] + self._additional_special_tokens = [] + + # (dsachan) Add BOS and EOS tokens + SPECIAL_TOKENS = {'eos_token': '[EOS]', + 'bos_token': '[BOS]'} + self._bos_token = '[BOS]' + self.add_token(self._bos_token) + self._bos_token_id = self.vocab.get(self._bos_token) + + self._eos_token = '[EOS]' + self.add_token(self._eos_token) + self._eos_token_id = self.vocab.get(self._eos_token) + + # (dsachan) Add additional special tokens + # These can be used as sentinel tokens in T5 model inputs + additional_special_tokens = [] + additional_special_tokens.extend( + ["".format(i) for i in range(vocab_extra_ids)]) + self.add_additional_special_tokens(additional_special_tokens) + + def add_token(self, token): + if token not in self.vocab: + self.inv_vocab[self.vocab_size] = token + # self.vocab_size comes from len(vocab) + # and it will increase as we add elements + self.vocab[token] = self.vocab_size + + def add_additional_special_tokens(self, tokens_list): + setattr(self, "additional_special_tokens", tokens_list) + for value in tokens_list: + self.add_token(value) + + @property + def vocab_size(self): + return self.tokenizer.vocab_size() + + @property + def vocab(self): + return self.tokenizer.vocab + + @property + def inv_vocab(self): + return self.tokenizer.inv_vocab + + def tokenize(self, text): + text_tokens = self.tokenizer.tokenize(text) + return self.tokenizer.convert_tokens_to_ids(text_tokens) + + def decode(self, ids): + tokens = self.tokenizer.convert_ids_to_tokens(ids) + return self.tokenizer.convert_tokens_to_string(tokens) + + def decode_token_ids(self, token_ids): + tokens = self.tokenizer.convert_ids_to_tokens(token_ids) + exclude_list = ['[PAD]', '[CLS]'] + non_pads = [t for t in tokens if t not in exclude_list] + + result = "" + for s in non_pads: + if s.startswith("##"): + result += s[2:] + else: + result += " " + s + + return result + + @property + def cls(self): + return self.cls_id + + @property + def sep(self): + return self.sep_id + + @property + def pad(self): + return self.pad_id + + @property + def mask(self): + return self.mask_id + + @property + def bos(self): + """ Id of the beginning of sentence token in the vocabulary.""" + return self._bos_token_id + + @property + def eos(self): + """ Id of the end of sentence token in the vocabulary.""" + return self._eos_token_id + + @property + def bos_token(self): + """ Beginning of sentence token id """ + return self._bos_token + + @property + def eos_token(self): + """ End of sentence token id """ + return self._eos_token + + @property + def additional_special_tokens(self): + """ All the additional special tokens you may want to use (list of strings).""" + return self._additional_special_tokens + + @property + def additional_special_tokens_ids(self): + """ Ids of all the additional special tokens in the vocabulary (list of integers).""" + return [self.vocab.get(token) for token in self._additional_special_tokens] + + @additional_special_tokens.setter + def additional_special_tokens(self, value): + self._additional_special_tokens = value + + +class _GPT2BPETokenizer(MegatronTokenizer): + """Original GPT2 BPE tokenizer.""" + + def __init__(self, vocab_file, merge_file): + super().__init__(vocab_file, merge_file) + + self.tokenizer = GPT2Tokenizer(vocab_file, merge_file, errors='replace', + special_tokens=[], max_len=None) + self.eod_id = self.tokenizer.encoder['<|endoftext|>'] + + @property + def vocab_size(self): + return len(self.tokenizer.encoder) + + @property + def vocab(self): + return self.tokenizer.encoder + + @property + def inv_vocab(self): + return self.tokenizer.decoder + + def tokenize(self, text): + return self.tokenizer.encode(text) + + def detokenize(self, token_ids): + return self.tokenizer.decode(token_ids) + + @property + def eod(self): + return self.eod_id + + +class _SentencePieceTokenizer(MegatronTokenizer): + """SentencePieceTokenizer-Megatron wrapper""" + + def __init__(self, model_file, vocab_extra_ids=0): + super().__init__(model_file, vocab_extra_ids=vocab_extra_ids) + + import sentencepiece + self.tokenizer = sentencepiece.SentencePieceProcessor(model_file=model_file) + self._initalize(vocab_extra_ids) + + def _populate_vocab(self): + self._vocab = {} + self._inv_vocab = {} + + for i in range(len(self.tokenizer)): + t = self.tokenizer.id_to_piece(i) + self._inv_vocab[i] = t + self._vocab[t] = i + + def _initalize(self, vocab_extra_ids): + self._populate_vocab() + self._special_tokens = {} + self._inv_special_tokens = {} + + self._t5_tokens = [] + + def _add_special_token(t): + if t not in self._vocab: + next_id = len(self._vocab) + self._vocab[t] = next_id + self._inv_vocab[next_id] = t + self._special_tokens[t] = self._vocab[t] + self._inv_special_tokens[self._vocab[t]] = t + + _add_special_token('') + self._cls_id = self._vocab[''] + _add_special_token('') + self._sep_id = self._vocab[''] + _add_special_token('') + self._eod_id = self._vocab[''] + _add_special_token('') + self._mask_id = self._vocab[''] + + pad_id = self.tokenizer.pad_id() + try: + pad_token = self.tokenizer.id_to_piece(pad_id) + except IndexError: + pad_token = '' + _add_special_token(pad_token) + self._pad_id = self._vocab[pad_token] + + bos_id = self.tokenizer.bos_id() + try: + bos_token = self.tokenizer.id_to_piece(bos_id) + except IndexError: + bos_token = '' + _add_special_token(bos_token) + self._bos_id = self._vocab[bos_token] + + eos_id = self.tokenizer.eos_id() + try: + eos_token = self.tokenizer.id_to_piece(eos_id) + except IndexError: + eos_token = '' + _add_special_token(eos_token) + self._eos_id = self._vocab[eos_token] + + for i in range(vocab_extra_ids): + t = "".format(i) + _add_special_token(t) + self._t5_tokens += [t] + + @property + def vocab_size(self): + return len(self._vocab) + + @property + def vocab(self): + return self._vocab + + @property + def inv_vocab(self): + return self._inv_vocab + + @property + def decoder(self): + return self._inv_vocab + + @property + def encoder(self): + return self._vocab + + # From: + # https://github.com/NVIDIA/NeMo/blob/c8fa217e811d60d11d014827c7f3845ff6c99ae7/nemo/collections/common/tokenizers/sentencepiece_tokenizer.py#L89 + def tokenize(self, text): + ids = [] + idx = 0 + + while 1: + indices = {} + for token in self._special_tokens: + try: + indices[token] = text[idx:].index(token) + except ValueError: + continue + if len(indices) == 0: + break + + next_token = min(indices, key=indices.get) + next_idx = idx + indices[next_token] + + ids.extend(self.tokenizer.encode_as_ids(text[idx:next_idx])) + ids.append(self._special_tokens[next_token]) + idx = next_idx + len(next_token) + + ids.extend(self.tokenizer.encode_as_ids(text[idx:])) + return ids + + # From: + # https://github.com/NVIDIA/NeMo/blob/c8fa217e811d60d11d014827c7f3845ff6c99ae7/nemo/collections/common/tokenizers/sentencepiece_tokenizer.py#L125 + def detokenize(self, ids): + text = "" + last_i = 0 + + for i, id in enumerate(ids): + if id in self._inv_special_tokens: + text += self.tokenizer.decode_ids(ids[last_i:i]) + " " + text += self._inv_special_tokens[id] + " " + last_i = i + 1 + + text += self.tokenizer.decode_ids(ids[last_i:]) + return text + + @property + def cls(self): + return self._cls_id + + @property + def sep(self): + return self._sep_id + + @property + def pad(self): + return self._pad_id + + @property + def bos(self): + return self._bos_id + + @property + def eod(self): + return self._eod_id + + @property + def eos(self): + return self._eos_id + + @property + def mask(self): + return self._mask_id + + @property + def additional_special_tokens_ids(self): + return [self.vocab[k] for k in self._t5_tokens] + + +class _GPTSentencePieceTokenizer(_SentencePieceTokenizer): + """SentencePieceTokenizer-Megatron wrapper""" + + def __init__(self, model_file,): + super().__init__(model_file, vocab_extra_ids=0) + + def _initalize(self, vocab_extra_ids): + self._populate_vocab() + + self._pad_id = self.tokenizer.pad_id() + self._bos_id = self.tokenizer.bos_id() + self._eos_id = self.tokenizer.eos_id() + + def tokenize(self, text): + return self.tokenizer.encode_as_ids(text) + + def detokenize(self, ids): + return self.tokenizer.decode_ids(ids) + + @property + def cls(self): + return -1 + + @property + def sep(self): + return -1 + + @property + def mask(self): + return -1 + + @property + def eod(self): + return self._eos_id + + @property + def additional_special_tokens_ids(self): + return None + + +class _Llama3Tokenizer(MegatronTokenizer): + def __init__(self,): + super().__init__("placeholder") + from transformers import AutoTokenizer + self.tokenizer = AutoTokenizer.from_pretrained( + "/data/llama3_8b_pretrain/llama3_8b_hf", + padding_side="right", + use_fast=False, + trust_remote_code=True + ) + self.extra_vocab_size = 256 + + @property + def vocab_size(self): + return self.tokenizer.vocab_size + self.extra_vocab_size + + @property + def vocab(self): + return self.tokenizer.encoder + + @property + def inv_vocab(self): + return self.tokenizer.decoder + + def tokenize(self, text): + return self.tokenizer.encode(text) + + def detokenize(self, token_ids): + return self.tokenizer.decode(token_ids) + + @property + def eod(self): + return self.tokenizer.eos_token_id + + @property + def eos_token(self): + return self.tokenizer.eos_token + + @property + def pad_token_id(self): + return self.tokenizer.pad_token_id + + +class _Llama2Tokenizer(_SentencePieceTokenizer): + """SentencePieceTokenizer-Megatron wrapper""" + + def __init__(self, model_file,): + super().__init__(model_file, vocab_extra_ids=0) + + def _initalize(self, vocab_extra_ids): + self._populate_vocab() + + # BOS / EOS token IDs + self.n_words: int = self.tokenizer.vocab_size() + self.bos_id: int = self.tokenizer.bos_id() + self.eos_id: int = self.tokenizer.eos_id() + self.pad_id: int = self.tokenizer.pad_id() + assert self.tokenizer.vocab_size() == self.tokenizer.get_piece_size() + + def tokenize(self, s: str, bos=True, eos=False): + '''Default args for text completion, not chat/dialog.''' + assert type(s) is str + t = self.tokenizer.encode(s) + if bos: + t = [self.bos_id] + t + if eos: + t = t + [self.eos_id] + return t + + def detokenize(self, ids): + return self.tokenizer.decode_ids(ids) + + @property + def cls(self): + return -1 + + @property + def sep(self): + return -1 + + @property + def mask(self): + return -1 + + @property + def eod(self): + return self.eos_id + + @property + def additional_special_tokens_ids(self): + return None + + +class _NullTokenizer: + def __init__(self, vocab_size): + vocab_size = int(vocab_size) + self._eos_id = vocab_size + self.vocab_size = vocab_size+1 + + def tokenize(self, text): + return [int(x) for x in text.split(' ')] + + def detokenize(self, ids): + text = [str(x) for x in ids] + return ' '.join(text) + + @property + def cls(self): + return -1 + + @property + def sep(self): + return -1 + + @property + def mask(self): + return -1 + + @property + def eod(self): + return self._eos_id + + @property + def additional_special_tokens_ids(self): + return None diff --git a/training/nvidia/docker_image/megatron_core060/Dockerfile b/training/nvidia/docker_image/megatron_core060/Dockerfile new file mode 100644 index 000000000..c5d46ae5a --- /dev/null +++ b/training/nvidia/docker_image/megatron_core060/Dockerfile @@ -0,0 +1,5 @@ +FROM nvcr.io/nvidia/pytorch:24.02-py3 +RUN /bin/bash -c "pip config set global.index-url https://mirror.baidu.com/pypi/simple" +RUN /bin/bash -c "uname -a" +RUN /bin/bash -c alias python3=python +RUN /bin/bash -c "git clone -b core_v0.6.0 https://github.com/NVIDIA/Megatron-LM.git" diff --git a/training/nvidia/docker_image/megatron_core060/megatron_core060_install.sh b/training/nvidia/docker_image/megatron_core060/megatron_core060_install.sh new file mode 100644 index 000000000..efca1a78d --- /dev/null +++ b/training/nvidia/docker_image/megatron_core060/megatron_core060_install.sh @@ -0,0 +1,3 @@ +#!/bin/bash +# using github mirrors to avoid github TTL + diff --git a/training/nvidia/llama3_8B-megatron/README.md b/training/nvidia/llama3_8B-megatron/README.md new file mode 100644 index 000000000..3cf5170f6 --- /dev/null +++ b/training/nvidia/llama3_8B-megatron/README.md @@ -0,0 +1,61 @@ + +### Nvidia GPU配置与运行信息参考 +#### 环境配置 +- ##### 硬件环境 + + - 机器型号: NVIDIA DGX A800(80G) + - 加速卡型号: NVIDIA_A800-SXM4-80GB + - CPU型号: AMD EPYC7742-64core@1.5G + - 多机网络类型、带宽: InfiniBand，200Gb/s + +- ##### 软件环境 + + - OS版本：Ubuntu 20.04 + - OS kernel版本: 5.4.0-126-generic + - 加速卡驱动版本：470.141.10 + - Docker 版本：20.10.18 + - 训练框架版本：megatron-core tag:core_v0.6.0 + - 依赖软件版本：sentencepiece==0.2.0, transformers==4.40.1 + +- ##### 并行策略 + + - 并行技术：张量、流水、数据混合并行，具体并行方案见“运行情况”章节 + - 实施者：megatron-core + - 实施细节：PP2DP4TP1 + +- ##### 优化策略 + + - flash attention 2 + - recompute-activations + - transformer-engine impl + +### 运行情况 + +* 输入批尺寸 + 1. local_batchsize(micro_batchsize)，简写为LBS，即实际进入模型的张量批尺寸，为config_A100x1x8.py中所写，在本case中默认为1。**厂商适配时可任意更改** + 2. seqlength(max_position_embedding)，简写为MPE，即实际进入模型的序列长度，为config_A100x1x8.py中所写，在本case中默认为8192，原则上不可更改 + 3. global_batchsize恒等于local_batchsize\*gradient_accumulate_steps\*data_parallel_size。在本case中，data_parallel_size=world_size/TPsize/PPsize。在本case中默认为512，使得globalbatchsize=4M tokens。 + +* 通用指标 + +| 指标名称 | 指标值 | 特殊说明 | +| ------- | --------------------- | ---------------------------------------- | +| 任务类别 | 自然语言理解 | | +| 模型 | llama3_8b | | +| 数据集 | wudao | wudao数据集来源于智源研究院
bin/idx数据集文件来源于阿里云灵骏团队
使用llama3 tokenizer预处理 | +| 数据精度 | precision,见“性能指标” | 可选fp32/amp/fp16/bf16 | +| 超参修改 | parallel,见“性能指标” | 格式为PPxDPyTPz，例如PP2DP4TP1 | +| 超参修改 | fix_hp,见“性能指标” | 跑满硬件设备评测吞吐量所需特殊超参 | +| 硬件设备简称 | nvidia A100 | | +| 硬件存储使用 | mem,见“性能指标” | 通常称为“显存”,单位为GiB | +| 计算使用率 | MFU,见“性能指标” | 参见PaLM论文定义 | +| **吞吐量** | **token/p/s,见“性能指标”** | 平均单卡每秒处理的token数 | + +* 性能指标 + +精度对齐需第21步及之后，所有步的loss与nvidia对应步的loss平均相对误差小于2%。NVloss曲线请联系智源研究院获取 + +| 配置 | precision | parallel | fix_hp | token/p/s | 是否精度对齐 | mem | MFU | +| ------------- | --------- | --------- | ------ | --------- | ------ | ----- | ----- | +| A100单机8卡（1x8） | bf16 | PP2DP4TP1 | / | 3698.7 | True | 76/80 | 56.9% | + diff --git a/training/nvidia/llama3_8B-megatron/config/config_A100x1x8.py b/training/nvidia/llama3_8B-megatron/config/config_A100x1x8.py new file mode 100644 index 000000000..9747f2131 --- /dev/null +++ b/training/nvidia/llama3_8B-megatron/config/config_A100x1x8.py @@ -0,0 +1,7 @@ +seqlength = 8192 +localbs = 1 +train_steps = 10 +theoryflops = 312000000000000.0 +epochs = 1 +tensor_parallel = 2 +pipeline_parallel = 1 diff --git a/training/nvidia/llama3_8B-megatron/config/requirements.txt b/training/nvidia/llama3_8B-megatron/config/requirements.txt new file mode 100644 index 000000000..20cdc820b --- /dev/null +++ b/training/nvidia/llama3_8B-megatron/config/requirements.txt @@ -0,0 +1,2 @@ +sentencepiece==0.2.0 +transformers==4.40.1 diff --git a/training/nvidia/llama3_8B-megatron/config/training_adapter.sh b/training/nvidia/llama3_8B-megatron/config/training_adapter.sh new file mode 100644 index 000000000..d95369427 --- /dev/null +++ b/training/nvidia/llama3_8B-megatron/config/training_adapter.sh @@ -0,0 +1 @@ +echo "[Prompt] nvidia adaption is NULL, for other Vendors" From 95cbba845a5cce966f9db2a7dd0ea382ea3f85f7 Mon Sep 17 00:00:00 2001 From: shh2000 <13820618441@163.com> Date: Sun, 28 Apr 2024 11:32:45 +0800 Subject: [PATCH 2/7] opt --- .../benchmarks/llama3_8B/megatron/README.md | 2 +- .../llama3_8B/megatron/arguments.patch | 7 + .../llama3_8B/megatron/arguments.py | 1615 ----------------- .../llama3_8B/megatron/pretrain_llama3.sh | 47 +- .../llama3_8B/megatron/run_pretraining.py | 52 +- .../llama3_8B/megatron/tokenizer.patch | 52 + .../llama3_8B/megatron/tokenizer.py | 567 ------ .../megatron_core060_install.sh | 2 - .../config/config_A100x1x8.py | 10 +- 9 files changed, 150 insertions(+), 2204 deletions(-) create mode 100644 training/benchmarks/llama3_8B/megatron/arguments.patch delete mode 100644 training/benchmarks/llama3_8B/megatron/arguments.py create mode 100644 training/benchmarks/llama3_8B/megatron/tokenizer.patch delete mode 100644 training/benchmarks/llama3_8B/megatron/tokenizer.py diff --git a/training/benchmarks/llama3_8B/megatron/README.md b/training/benchmarks/llama3_8B/megatron/README.md index 1f30bc795..ed07a6ef4 100644 --- a/training/benchmarks/llama3_8B/megatron/README.md +++ b/training/benchmarks/llama3_8B/megatron/README.md @@ -15,7 +15,7 @@ This case includes code from the LLAMA 3 COMMUNITY LICENSE AGREEMENT License ope ### 模型配置及tokenizer准备 -本测试样例为预训练case，需要下载tokenizer，不需要下载模型代码或模型权重。tokenizer需向llama3官方申请并下载8B版本，并在data_dir下创建llama3_8b_hf目录，按照[unsloth/llama-3-8b at main (huggingface.co)](https://huggingface.co/unsloth/llama-3-8b/tree/main)仓库示例格式存放（不需要权重部分）。 +本测试样例为预训练case，需要下载tokenizer，不需要下载模型代码或模型权重。tokenizer需向llama3官方申请并下载8B版本，并在data_dir下创建\目录，按照[unsloth/llama-3-8b at main (huggingface.co)](https://huggingface.co/unsloth/llama-3-8b/tree/main)仓库示例格式存放（不需要权重部分）。其中，\目录可更改，需在形如config\_A100\_1x8.py的配置文件中同步修改。默认为\。 ### 数据集准备 diff --git a/training/benchmarks/llama3_8B/megatron/arguments.patch b/training/benchmarks/llama3_8B/megatron/arguments.patch new file mode 100644 index 000000000..9bda43d3d --- /dev/null +++ b/training/benchmarks/llama3_8B/megatron/arguments.patch @@ -0,0 +1,7 @@ +# Copyright (c) 2024 BAAI. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License") +#!/usr/bin/env python3 +# -*- coding: UTF-8 -*- +1411a1412 +> 'Llama3Tokenizer', diff --git a/training/benchmarks/llama3_8B/megatron/arguments.py b/training/benchmarks/llama3_8B/megatron/arguments.py deleted file mode 100644 index 5a062956f..000000000 --- a/training/benchmarks/llama3_8B/megatron/arguments.py +++ /dev/null @@ -1,1615 +0,0 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. - -"""Megatron arguments.""" - -import argparse -import dataclasses -import json -import os -import torch -import types - -import torch.nn.functional as F -from megatron.core.models.retro.utils import ( - get_config_path as get_retro_config_path, - get_gpt_data_dir as get_retro_data_dir, -) -from megatron.core.transformer import TransformerConfig - - -def parse_args(extra_args_provider=None, ignore_unknown_args=False): - """Parse all arguments.""" - parser = argparse.ArgumentParser(description='Megatron-LM Arguments', - allow_abbrev=False) - - # Standard arguments. - parser = _add_network_size_args(parser) - parser = _add_regularization_args(parser) - parser = _add_training_args(parser) - parser = _add_initialization_args(parser) - parser = _add_learning_rate_args(parser) - parser = _add_checkpointing_args(parser) - parser = _add_mixed_precision_args(parser) - parser = _add_distributed_args(parser) - parser = _add_validation_args(parser) - parser = _add_data_args(parser) - parser = _add_autoresume_args(parser) - parser = _add_biencoder_args(parser) - parser = _add_vision_args(parser) - parser = _add_moe_args(parser) - parser = _add_logging_args(parser) - parser = _add_inference_args(parser) - parser = _add_transformer_engine_args(parser) - parser = _add_retro_args(parser) - parser = _add_experimental_args(parser) - - # Custom arguments. - if extra_args_provider is not None: - parser = extra_args_provider(parser) - - # Parse. - if ignore_unknown_args: - args, _ = parser.parse_known_args() - else: - args = parser.parse_args() - - # Experimental yaml - if args.yaml_cfg is not None: - from .yaml_arguments import load_yaml - assert args.yaml_cfg and args.use_mcore_models, "To use yaml, mcore must be enabled" - args = load_yaml(args.yaml_cfg) - - - # Args from environment - args.rank = int(os.getenv('RANK', '0')) - args.world_size = int(os.getenv("WORLD_SIZE", '1')) - - return args - - -def load_retro_config(retro_project_dir): - '''Load Retro's config.json.''' - - # Retro config path. - retro_config_path = get_retro_config_path(retro_project_dir) - assert os.path.exists(retro_config_path), \ - "Retro project dir missing config.json." - - # Load retro config. - with open(retro_config_path) as f: - retro_config = types.SimpleNamespace(**json.load(f)) - - return retro_config - - -def load_retro_args(args): - """Load predefined args from Retro config (if applicable). - - When using Retro (or GPT for comparison purposes), data arguments are - overridden by the saved config.json within the Retro project directory. This - is to ensure that the data used for pretraining is consistent with the data - that was preprocessed using the Retro preprocessing pipeline (see - `tools/retro/preprocess_data.py`). - """ - - # Return if no project directory is specified. - if args.retro_project_dir is None: - return - - # Load retro config. - retro_config = load_retro_config(args.retro_project_dir) - - # Retro data path is relative to project dir (via hard or soft links). - data_dir = get_retro_data_dir(args.retro_project_dir) - data_path = list(retro_config.retro_gpt_data_path) - if len(data_path) % 2 == 0: - for i in range(len(data_path) - 1, -1, -2): - data_path[i] = os.path.join(data_dir, data_path[i]) - else: - assert len(data_path) == 1 - data_path[0] = os.path.join(data_dir, data_path[0]) - - # Update args. - args.data_cache_path = retro_config.retro_gpt_data_cache_path - args.data_path = data_path if args.data_path is None else args.data_path - args.eval_interval = retro_config.retro_gpt_eval_interval - args.eval_iters = retro_config.retro_gpt_eval_iters - args.global_batch_size = retro_config.retro_gpt_global_batch_size - args.max_position_embeddings = retro_config.retro_gpt_seq_length - args.merge_file = os.path.join( - args.retro_project_dir, - retro_config.retro_gpt_merge_file, - ) if retro_config.retro_gpt_merge_file is not None else None - args.seed = retro_config.retro_gpt_seed - args.seq_length = retro_config.retro_gpt_seq_length - args.tokenizer_model = os.path.join( - args.retro_project_dir, - retro_config.retro_gpt_tokenizer_model, - ) if retro_config.retro_gpt_tokenizer_model is not None else None - args.tokenizer_type = retro_config.retro_gpt_tokenizer_type - args.train_samples = retro_config.retro_gpt_train_samples - args.vocab_file = os.path.join( - args.retro_project_dir, - retro_config.retro_gpt_vocab_file, - ) if retro_config.retro_gpt_vocab_file is not None else None - - # Retro-specific args. - args.retro_block_size = retro_config.retro_block_size - args.retro_chunk_length = retro_config.retro_gpt_chunk_length - args.retro_neighbor_dirs = retro_config.retro_neighbor_dirs - args.retro_split_preprocessing = retro_config.retro_gpt_split - args.retro_bert_tokenizer_type = retro_config.retro_bert_tokenizer_type - args.retro_bert_vocab_file = retro_config.retro_bert_vocab_file - - -def validate_args(args, defaults={}): - - # Load saved args from Retro (if applicable). - load_retro_args(args) - - # Tensor model parallel size. - args.tensor_model_parallel_size = min( - args.tensor_model_parallel_size, args.world_size) - assert args.world_size % args.tensor_model_parallel_size == 0, 'world size'\ - ' ({}) is not divisible by tensor model parallel size ({})'.format( - args.world_size, args.tensor_model_parallel_size) - - # Pipeline model parallel size. - args.pipeline_model_parallel_size = min( - args.pipeline_model_parallel_size, - (args.world_size // args.tensor_model_parallel_size)) - args.transformer_pipeline_model_parallel_size = ( - args.pipeline_model_parallel_size - 1 - if args.standalone_embedding_stage else - args.pipeline_model_parallel_size - ) - - # Checks. - model_parallel_size = args.pipeline_model_parallel_size * \ - args.tensor_model_parallel_size - assert args.world_size % (model_parallel_size * args.context_parallel_size) == 0, \ - 'world size ({}) is not divisible by tensor parallel size ({}) times ' \ - 'pipeline parallel size ({}) times context parallel size ({})'.format( - args.world_size, args.tensor_model_parallel_size, - args.pipeline_model_parallel_size, args.context_parallel_size) - args.data_parallel_size = args.world_size // (model_parallel_size * args.context_parallel_size) - if args.rank == 0: - print('using world size: {}, data-parallel size: {}, ' - 'context-parallel size: {} ' - 'tensor-model-parallel size: {}, ' - 'pipeline-model-parallel size: {} '.format( - args.world_size, args.data_parallel_size, - args.context_parallel_size, - args.tensor_model_parallel_size, - args.pipeline_model_parallel_size), flush=True) - if args.pipeline_model_parallel_size > 1: - if args.pipeline_model_parallel_split_rank is not None: - assert args.pipeline_model_parallel_split_rank < \ - args.pipeline_model_parallel_size, 'split rank needs'\ - ' to be less than pipeline model parallel size ({})'.format( - args.pipeline_model_parallel_size) - - if args.tp_comm_overlap: - assert args.sequence_parallel == True, 'Tensor parallel communication/GEMM overlap can happen only when sequence parallelism is enabled' - - # Deprecated arguments - assert args.batch_size is None, '--batch-size argument is no longer ' \ - 'valid, use --micro-batch-size instead' - del args.batch_size - assert args.warmup is None, '--warmup argument is no longer valid, use ' \ - '--lr-warmup-fraction instead' - del args.warmup - assert args.model_parallel_size is None, '--model-parallel-size is no ' \ - 'longer valid, use --tensor-model-parallel-size instead' - del args.model_parallel_size - - if args.checkpoint_activations: - if args.rank == 0: - print('--checkpoint-activations is no longer valid, use --recompute-activations, ' - 'or, for more control, --recompute-granularity and --recompute-method.') - exit() - del args.checkpoint_activations - - if args.recompute_activations: - args.recompute_granularity = 'selective' - del args.recompute_activations - - # Set input defaults. - for key in defaults: - # For default to be valid, it should not be provided in the - # arguments that are passed to the program. We check this by - # ensuring the arg is set to None. - if getattr(args, key, None) is not None: - if args.rank == 0: - print('WARNING: overriding default arguments for {key}:{v} \ - with {key}:{v2}'.format(key=key, v=defaults[key], - v2=getattr(args, key)), - flush=True) - else: - setattr(args, key, defaults[key]) - - # Batch size. - assert args.micro_batch_size is not None - assert args.micro_batch_size > 0 - if args.global_batch_size is None: - args.global_batch_size = args.micro_batch_size * args.data_parallel_size - if args.rank == 0: - print('setting global batch size to {}'.format( - args.global_batch_size), flush=True) - assert args.global_batch_size > 0 - if args.num_layers_per_virtual_pipeline_stage is not None: - assert args.pipeline_model_parallel_size > 2, \ - 'pipeline-model-parallel size should be greater than 2 with ' \ - 'interleaved schedule' - assert args.num_layers % args.transformer_pipeline_model_parallel_size == 0, \ - 'number of layers should be divisible by the pipeline parallel size' - num_layers_per_pipeline_stage = args.num_layers // args.transformer_pipeline_model_parallel_size - assert num_layers_per_pipeline_stage % args.num_layers_per_virtual_pipeline_stage == 0, \ - 'number of layers per pipeline stage must be divisible number of layers per virtual pipeline stage' - args.virtual_pipeline_model_parallel_size = num_layers_per_pipeline_stage // \ - args.num_layers_per_virtual_pipeline_stage - else: - args.virtual_pipeline_model_parallel_size = None - # Overlap P2P communication is disabled if not using the interleaved schedule. - args.overlap_p2p_comm = False - if args.rank == 0: - print('WARNING: Setting args.overlap_p2p_comm to False since non-interleaved ' - 'schedule does not support overlapping p2p communication') - - if args.overlap_param_gather: - assert args.use_distributed_optimizer, \ - '--overlap-param-gather only supported with distributed optimizer' - assert args.overlap_grad_reduce, \ - '--overlap-grad-reduce should be turned on when using --overlap-param-gather' - assert args.use_mcore_models, \ - '--overlap-param-gather only supported with MCore models' - - # Parameters dtype. - args.params_dtype = torch.float - if args.fp16: - assert not args.bf16 - args.params_dtype = torch.half - # Turn off checking for NaNs in loss and grads if using dynamic loss scaling, - # where NaNs in grads / loss are signal to the loss scaler. - if not args.loss_scale: - args.check_for_nan_in_loss_and_grad = False - if args.rank == 0: - print('WARNING: Setting args.check_for_nan_in_loss_and_grad to False since ' - 'dynamic loss scaling is being used') - if args.bf16: - assert not args.fp16 - args.params_dtype = torch.bfloat16 - # bfloat16 requires gradient accumulation and all-reduce to - # be done in fp32. - if not args.accumulate_allreduce_grads_in_fp32: - args.accumulate_allreduce_grads_in_fp32 = True - if args.rank == 0: - print('accumulate and all-reduce gradients in fp32 for ' - 'bfloat16 data type.', flush=True) - - if args.rank == 0: - print('using {} for parameters ...'.format(args.params_dtype), - flush=True) - - if args.dataloader_type is None: - args.dataloader_type = 'single' - - # Consumed tokens. - args.consumed_train_samples = 0 - args.consumed_valid_samples = 0 - - # Support for variable sequence lengths across batches/microbatches. - # set it if the dataloader supports generation of variable sequence lengths - # across batches/microbatches. Due to additional communication overhead - # during pipeline parallelism, it should not be set if sequence length - # is constant during training. - args.variable_seq_lengths = False - - # Iteration-based training. - if args.train_iters: - # If we use iteration-based training, make sure the - # sample-based options are off. - assert args.train_samples is None, \ - 'expected iteration-based training' - assert args.lr_decay_samples is None, \ - 'expected iteration-based learning rate decay' - assert args.lr_warmup_samples == 0, \ - 'expected iteration-based learning rate warmup' - assert args.rampup_batch_size is None, \ - 'expected no batch-size rampup for iteration-based training' - if args.lr_warmup_fraction is not None: - assert args.lr_warmup_iters == 0, \ - 'can only specify one of lr-warmup-fraction and lr-warmup-iters' - - # Sample-based training. - if args.train_samples: - # If we use sample-based training, make sure the - # iteration-based options are off. - assert args.train_iters is None, \ - 'expected sample-based training' - assert args.lr_decay_iters is None, \ - 'expected sample-based learning rate decay' - assert args.lr_warmup_iters == 0, \ - 'expected sample-based learnig rate warmup' - if args.lr_warmup_fraction is not None: - assert args.lr_warmup_samples == 0, \ - 'can only specify one of lr-warmup-fraction ' \ - 'and lr-warmup-samples' - - if args.num_layers is not None: - assert args.encoder_num_layers is None, \ - 'cannot have both num-layers and encoder-num-layers specified' - args.encoder_num_layers = args.num_layers - else: - assert args.encoder_num_layers is not None, \ - 'either num-layers or encoder-num-layers should be specified' - args.num_layers = args.encoder_num_layers - - # Check required arguments. - required_args = ['num_layers', 'hidden_size', 'num_attention_heads', - 'max_position_embeddings'] - for req_arg in required_args: - _check_arg_is_not_none(args, req_arg) - - # Checks. - if args.ffn_hidden_size is None: - if args.swiglu: - # reduce the dimnesion for MLP since projections happens on - # two linear layers. this keeps the number of paramters in - # the same ballpark as the counterpart with 4*h size - # we keep it a multiple of 64, which means the actual tensor size - # will be a multiple of 64 / tp_size - args.ffn_hidden_size = int((4 * args.hidden_size * 2 / 3) / 64) * 64 - else: - args.ffn_hidden_size = 4 * args.hidden_size - - if args.kv_channels is None: - assert args.hidden_size % args.num_attention_heads == 0 - args.kv_channels = args.hidden_size // args.num_attention_heads - - if args.seq_length is not None: - assert args.encoder_seq_length is None - args.encoder_seq_length = args.seq_length - else: - assert args.encoder_seq_length is not None - args.seq_length = args.encoder_seq_length - - if args.seq_length is not None: - assert args.max_position_embeddings >= args.seq_length - if args.decoder_seq_length is not None: - assert args.max_position_embeddings >= args.decoder_seq_length - if args.lr is not None: - assert args.min_lr <= args.lr - if args.save is not None: - assert args.save_interval is not None - # Mixed precision checks. - if args.fp16_lm_cross_entropy: - assert args.fp16, 'lm cross entropy in fp16 only support in fp16 mode.' - if args.fp32_residual_connection: - assert args.fp16 or args.bf16, \ - 'residual connection in fp32 only supported when using fp16 or bf16.' - - if args.moe_grouped_gemm: - assert args.bf16, 'Currently GroupedGEMM for MoE only supports bf16 dtype.' - dc = torch.cuda.get_device_capability() - assert dc[0] >= 8, "Unsupported compute capability for GroupedGEMM kernels." - - if args.weight_decay_incr_style == 'constant': - assert args.start_weight_decay is None - assert args.end_weight_decay is None - args.start_weight_decay = args.weight_decay - args.end_weight_decay = args.weight_decay - else: - assert args.start_weight_decay is not None - assert args.end_weight_decay is not None - - TORCH_MAJOR = int(torch.__version__.split('.')[0]) - TORCH_MINOR = int(torch.__version__.split('.')[1]) - # Persistent fused layer norm. - if TORCH_MAJOR < 1 or (TORCH_MAJOR == 1 and TORCH_MINOR < 11): - args.no_persist_layer_norm = True - if args.rank == 0: - print('Persistent fused layer norm kernel is supported from ' - 'pytorch v1.11 (nvidia pytorch container paired with v1.11). ' - 'Defaulting to no_persist_layer_norm=True') - - # Activation recomputing. - if args.distribute_saved_activations: - assert args.tensor_model_parallel_size > 1, 'can distribute ' \ - 'recomputed activations only across tensor model ' \ - 'parallel groups' - assert args.recompute_granularity == 'full', \ - 'distributed recompute activations is only '\ - 'application to full recompute granularity' - assert args.recompute_method is not None, \ - 'for distributed recompute activations to work you '\ - 'need to use a recompute method ' - assert (TORCH_MAJOR, TORCH_MINOR) >= (1, 10), \ - 'distributed recompute activations are supported for pytorch ' \ - 'v1.10 and above (Nvidia Pytorch container >= 21.07). Current ' \ - 'pytorch version is v%s.%s.' % (TORCH_MAJOR, TORCH_MINOR) - - if args.recompute_granularity == 'selective': - assert args.recompute_method is None, \ - 'recompute method is not yet supported for ' \ - 'selective recomputing granularity' - - # disable sequence parallelism when tp=1 - # to avoid change in numerics when - # sequence_parallelism is enabled. - if args.tensor_model_parallel_size == 1: - args.sequence_parallel = False - - # disable async_tensor_model_parallel_allreduce when - # model parallel memory optimization is enabled - if args.sequence_parallel: - args.async_tensor_model_parallel_allreduce = False - - if os.environ.get('CUDA_DEVICE_MAX_CONNECTIONS') != "1": - if args.sequence_parallel: - raise RuntimeError( - "Using sequence parallelism requires setting the environment variable " - "CUDA_DEVICE_MAX_CONNECTIONS to 1") - if args.async_tensor_model_parallel_allreduce: - raise RuntimeError( - "Using async gradient all reduce requires setting the environment " - "variable CUDA_DEVICE_MAX_CONNECTIONS to 1") - - # Disable bias gelu fusion if we are disabling bias altogether - if not args.add_bias_linear: - args.bias_gelu_fusion = False - - # Retro checks. - if args.retro_add_retriever: - - # Train samples should be auto-loaded. - assert args.train_samples is not None, \ - "args.train_samples should be auto-loaded from the retro config." - - # Sequence parallelism unsupported. - assert not args.sequence_parallel, \ - "retro currently does not support sequence parallelism." - - # Pipeline parallelism unsupported. - assert args.pipeline_model_parallel_size == 1, \ - "retro currently does not support pipeline parallelism." - - if args.decoupled_lr is not None or args.decoupled_min_lr is not None: - assert args.use_mcore_models, \ - '--decoupled-lr and --decoupled-min-lr only supported by Megatron Core, please add --use-mcore-models.' - - # Legacy RoPE arguments - if args.use_rotary_position_embeddings: - args.position_embedding_type = 'rope' - if args.rotary_interleaved and args.apply_rope_fusion: - raise RuntimeError('--rotary-interleaved does not work with rope_fusion.') - if args.rotary_interleaved and not args.use_mcore_models: - raise RuntimeError('--rotary-interleaved only support Megatron Core, please add --use-mcore-models.') - - # Would just need to add 'NoPE' as a position_embedding_type to support this, but for now - # don't allow it to keep things simple - if not args.add_position_embedding and args.position_embedding_type != 'rope': - raise RuntimeError('--no-position-embedding is deprecated, use --position-embedding-type') - - # MoE Spec check - if args.num_experts is not None: - assert args.spec is None, "Model Spec must be None when using MoEs" - if args.tensor_model_parallel_size > 1: - assert args.sequence_parallel, \ - "When using MoE and tensor parallelism, sequence parallelism must be used." - - # Expert parallelism check - if args.expert_model_parallel_size > 1: - assert args.num_experts is not None, "num_experts must be non None to use expert model parallelism" - assert args.num_experts % args.expert_model_parallel_size == 0, \ - "Number of experts should be a multiple of expert model parallel_size." - assert not args.fp16, \ - "Expert parallelism is not supported with fp16 training." - - # Distributed checkpointing checks - if args.use_dist_ckpt and not args.use_mcore_models: - raise RuntimeError('--use-dist-ckpt only support Megatron Core, please add --use-mcore-models.') - - # Print arguments. - _print_args("arguments", args) - - return args - - -def _print_args(title, args): - """Print arguments.""" - if args.rank == 0: - print(f'------------------------ {title} ------------------------', - flush=True) - str_list = [] - for arg in vars(args): - dots = '.' * (48 - len(arg)) - str_list.append(' {} {} {}'.format(arg, dots, getattr(args, arg))) - for arg in sorted(str_list, key=lambda x: x.lower()): - print(arg, flush=True) - print(f'-------------------- end of {title} ---------------------', - flush=True) - - -def _check_arg_is_not_none(args, arg): - assert getattr(args, arg) is not None, '{} argument is None'.format(arg) - - -def core_transformer_config_from_args(args, config_class=None): - - # Config class. - config_class = config_class or TransformerConfig - - # Translate args to core transformer configuration - kw_args = {} - for f in dataclasses.fields(config_class): - if hasattr(args, f.name): - kw_args[f.name] = getattr(args, f.name) - kw_args['persist_layer_norm'] = not args.no_persist_layer_norm - kw_args['layernorm_zero_centered_gamma'] = args.apply_layernorm_1p - kw_args['layernorm_epsilon'] = args.norm_epsilon - kw_args['deallocate_pipeline_outputs'] = True - kw_args['pipeline_dtype'] = args.params_dtype - kw_args['batch_p2p_comm'] = not args.overlap_p2p_comm - kw_args['num_moe_experts'] = args.num_experts - kw_args['rotary_interleaved'] = args.rotary_interleaved - if args.swiglu: - kw_args['activation_func'] = F.silu - kw_args['gated_linear_unit'] = True - kw_args['bias_activation_fusion'] = args.bias_swiglu_fusion - else: - kw_args['bias_activation_fusion'] = args.bias_gelu_fusion - if args.squared_relu: - assert not args.swiglu - def squared_relu(x): - return torch.pow(F.relu(x), 2) - kw_args['activation_func'] = squared_relu - if args.init_method_xavier_uniform: - kw_args['init_method'] = torch.nn.init.xavier_uniform_ - kw_args['scaled_init_method'] = torch.nn.init.xavier_uniform_ - if args.group_query_attention: - kw_args['num_query_groups'] = args.num_query_groups - else: - kw_args['num_query_groups'] = None - - # Return config. - return config_class(**kw_args) - - -def _add_transformer_engine_args(parser): - group = parser.add_argument_group(title='Transformer-Engine') - - group.add_argument('--fp8-format', default=None, - choices=['e4m3', 'hybrid'], - help='Which fp8 format scheme to use for FP8 tensors in the forward and backward pass', - dest='fp8') - group.add_argument('--fp8-margin', type=int, default=0, - help='Scaling margin for fp8', - dest='fp8_margin') - group.add_argument('--fp8-interval', type=int, default=1, - help='Scaling update interval for fp8', - dest='fp8_interval') - group.add_argument('--fp8-amax-history-len', type=int, default=1, - help='Number of steps for which amax history is recorded per tensor', - dest='fp8_amax_history_len') - group.add_argument('--fp8-amax-compute-algo', default='most_recent', - choices=['most_recent', 'max'], - help='Algorithm for computing amax from history', - dest='fp8_amax_compute_algo') - group.add_argument('--no-fp8-wgrad', action='store_false', - help='Execute wgrad in higher precision even for FP8 runs', - dest='fp8_wgrad') - group.add_argument('--transformer-impl', default='transformer_engine', - choices=['local', 'transformer_engine'], - help='Which Transformer implementation to use.') - - return parser - -def _add_inference_args(parser): - group = parser.add_argument_group(title='inference') - - group.add_argument('--inference-batch-times-seqlen-threshold', - type=int, default=512, - help='During inference, if batch-size times ' - 'sequence-length is smaller than this threshold ' - 'then we will not use pipelining, otherwise we will.') - group.add_argument('--max-tokens-to-oom', - type=int, default=12000, - help='Maximum number of tokens during inference' - 'tokens here is # in prompt + # to generate' - 'Allows us to throw an error before OOM crashes server') - group.add_argument('--output-bert-embeddings', action='store_true', - help='Output Bert embeddings (via mean pooling) from ' - 'model, rather than its binary head output or entire ' - 'hidden batch.') - group.add_argument('--bert-embedder-type', default="megatron", - choices=["megatron", "huggingface"], - help='Select either Megatron or Huggingface as the ' - 'Bert embedder.') - - return parser - - -def _add_retro_args(parser): - group = parser.add_argument_group(title='retro') - - group.add_argument('--retro-project-dir', default=None, - help='Retro project directory, which contains the ' - 'preprocessed data for pretraining. This directory ' - 'is built during preprocessing (see ' - 'tools/retro/README.md), and contains subdirectories ' - 'for the chunk database and pretraining neighbors.') - group.add_argument('--retro-add-retriever', - action='store_true', default=False, - help='Add a retriever to the transformer, for use in ' - 'pretraining a Retro model.') - group.add_argument('--retro-cyclic-train-iters', type=int, default=None, - help='Set number of training iterations for cyclic ' - 'Retro training.') - group.add_argument('--retro-encoder-layers', type=int, default=2, - help='Number of layers to use for the retrieval ' - 'encoder.') - group.add_argument('--retro-encoder-hidden-dropout', - type=float, default=0.1, help='Hidden dropout for ' - 'retrieval encoder.') - group.add_argument('--retro-encoder-attention-dropout', - type=float, default=0.1, help='Attention dropout for ' - 'retrieval encoder.') - group.add_argument("--retro-num-neighbors", type=int, default=2, - help='Number of neighbors to retrieve during ' - 'pretraining.') - group.add_argument("--retro-num-retrieved-chunks", type=int, default=2, - help='Number of chunks to retrieve from the retrieval ' - 'database.') - group.add_argument("--retro-attention-gate", type=float, default=1, - help="Gated cross attention.") - group.add_argument("--retro-no-verify-neighbor-count", action="store_false", - dest="retro_verify_neighbor_count", - help="Skip verifying that len(GPT dataset) == len(saved " - "neighbors).") - - # Enforce argument naming convention. - for action in group._group_actions: - prefix = action.dest.split("_")[0] - assert prefix == "retro", \ - "Retro args must be prefixed with '--retro-*', for consistent " \ - "styling. Please fix '%s'." % ", ".join(action.option_strings) - - return parser - - -def _add_network_size_args(parser): - group = parser.add_argument_group(title='network size') - - group.add_argument('--num-layers', type=int, default=None, - help='Number of transformer layers.') - group.add_argument('--encoder-num-layers', type=int, default=None, - help='Number of encoder transformer layers.') - group.add_argument('--decoder-num-layers', type=int, default=None, - help='Number of decoder transformer layers.') - group.add_argument('--hidden-size', type=int, default=None, - help='Tansformer hidden size.') - group.add_argument('--ffn-hidden-size', type=int, default=None, - help='Transformer Feed-Forward Network hidden size. ' - 'This is set to 4*hidden-size if not provided') - group.add_argument('--num-attention-heads', type=int, default=None, - help='Number of transformer attention heads.') - group.add_argument('--kv-channels', type=int, default=None, - help='Projection weights dimension in multi-head ' - 'attention. This is set to ' - ' args.hidden_size // args.num_attention_heads ' - 'if not provided.') - group.add_argument('--group-query-attention', action='store_true', - help='Use group-query attention.') - group.add_argument('--num-query-groups', type=int, default=1) - - group.add_argument('--max-position-embeddings', type=int, default=None, - help='Maximum number of position embeddings to use. ' - 'This is the size of position embedding.') - group.add_argument('--position-embedding-type', type=str, default='learned_absolute', - choices=['learned_absolute', 'rope'], - help='Position embedding type.') - group.add_argument('--use-rotary-position-embeddings', action='store_true', - help='Use rotary positional embeddings or not. ' - 'Deprecated: use --position-embedding-type') - group.add_argument('--rotary-percent', type=float, default=1.0, - help='Percent of rotary dimension to use, default 100%%') - group.add_argument('--rotary-interleaved', action='store_true', - help='Use interleaved rotary embedding.') - group.add_argument('--rotary-seq-len-interpolation-factor', type=int, default=None, - help='Sequence length interpolation factor for rotary embeddings.') - group.add_argument('--no-position-embedding', - action='store_false', - help='Disable position embedding. Deprecated: use --position-embedding-type', - dest='add_position_embedding') - group.add_argument('--make-vocab-size-divisible-by', type=int, default=128, - help='Pad the vocab size to be divisible by this value.' - 'This is added for computational efficieny reasons.') - group.add_argument('--normalization', default='LayerNorm', - choices=['LayerNorm', 'RMSNorm'], - help='Which normalization technique to use.') - group.add_argument('--norm-epsilon', type=float, default=1e-5, - help='Epsilon for layer norm and RMS norm.') - group.add_argument('--apply-layernorm-1p', action='store_true', - help='Adjust LayerNorm weights such that they are centered ' - 'around zero. This improves numerical stability.') - group.add_argument('--apply-residual-connection-post-layernorm', - action='store_true', - help='If set, use original BERT residula connection ' - 'ordering.') - group.add_argument('--openai-gelu', action='store_true', - help='Use OpenAIs GeLU implementation. This option' - 'should not be used unless for backward compatibility' - 'reasons.') - group.add_argument('--squared-relu', action='store_true', - help='Use squared relu activation instead of default gelu') - group.add_argument('--swiglu', action='store_true', - help='Use gated linear units and SiLU activation instead of default gelu') - group.add_argument('--onnx-safe', type=bool, required=False, - help='Use workarounds for known problems with ' - 'Torch ONNX exporter') - group.add_argument('--bert-no-binary-head', action='store_false', - help='Disable BERT binary head.', - dest='bert_binary_head') - group.add_argument('--untie-embeddings-and-output-weights', action='store_true', - help='Untie embeddings and output weights.'), - return parser - - -def _add_logging_args(parser): - group = parser.add_argument_group(title='logging') - - group.add_argument('--log-params-norm', action='store_true', - help='If set, calculate and log parameters norm.') - group.add_argument('--log-num-zeros-in-grad', action='store_true', - help='If set, calculate and log the number of zeros in gradient.') - group.add_argument('--log-throughput', action='store_true', - help='If set, calculate and log throughput per GPU.') - group.add_argument('--log-progress', action='store_true', - help='If set, log progress (in terms of number of processed tokens and ' - 'number of floating-point operations) to progress.txt file in checkpoint ' - 'directory.') - group.add_argument('--timing-log-level', type=int, - default=0, choices=range(0,3), - help='Granularity level to measure and report timing. ' - ' 0: report only iteration time and make sure timing ' - ' does not introduce extra overhead.' - ' 1: report timing for operations that are executed ' - ' very limited times (basically once) during ' - ' each iteration (such as gradient all-reduce) ' - ' 2: report timing for operations that migh be ' - ' executed numerous times during each iteration. ' - 'Note that setting the level to 1 or 2 might ' - 'cause increase in iteration time.') - group.add_argument('--no-barrier-with-level-1-timing', action='store_false', - help='If not set, use barrier with level 1 time ' - 'measurements. Note that this is up to the user ' - 'to make sure calling barrier with their timers ' - 'will not result in hangs. This can happen if for ' - 'example the user adds a level 1 timer that is not ' - 'called by all ranks.', - dest='barrier_with_L1_time') - group.add_argument('--timing-log-option', type=str, default='minmax', - choices=['max', 'minmax', 'all'], - help='Options for logging timing:' - ' max: report the max timing across all ranks' - ' minmax: report min and max timings across all ranks' - ' all: report timings of all ranks.') - group.add_argument('--tensorboard-log-interval', type=int, default=1, - help='Report to tensorboard interval.') - group.add_argument('--tensorboard-queue-size', type=int, default=1000, - help='Size of the tensorboard queue for pending events ' - 'and summaries before one of the ‘add’ calls forces a ' - 'flush to disk.') - group.add_argument('--log-timers-to-tensorboard', action='store_true', - help='If set, write timers to tensorboard.') - group.add_argument('--log-batch-size-to-tensorboard', action='store_true', - help='If set, write batch-size to tensorboard.') - group.add_argument('--no-log-learnig-rate-to-tensorboard', - action='store_false', - help='Disable learning rate logging to tensorboard.', - dest='log_learning_rate_to_tensorboard') - group.add_argument('--no-log-loss-scale-to-tensorboard', - action='store_false', - help='Disable loss-scale logging to tensorboard.', - dest='log_loss_scale_to_tensorboard') - group.add_argument('--log-validation-ppl-to-tensorboard', - action='store_true', - help='If set, write validation perplexity to ' - 'tensorboard.') - group.add_argument('--log-memory-to-tensorboard', - action='store_true', - help='Enable memory logging to tensorboard.') - group.add_argument('--log-world-size-to-tensorboard', - action='store_true', - help='Enable world size logging to tensorboard.') - group.add_argument('--wandb-project', type=str, default='', - help='The wandb project name. Ignore wandb by default.') - group.add_argument('--wandb-exp-name', type=str, default='', - help='The wandb experiment name.') - group.add_argument('--wandb-save-dir', type=str, default='', - help='Path to save the wandb results locally.') - group.add_argument('--enable-one-logger', action='store_true', - help='If set, use one_logger to track E2E metrics' - 'Note that one_logger is an internal tool and not available externally. ' - 'For installation, please try command: `pip install ' - '--index-url=https://sc-hw-artf.nvidia.com/api/pypi/hwinf-ml-pypi/simple' - ' one_logger` or go to https://gitlab-master.nvidia.com/hwinf-dcm/onelogger ' - 'for more details') - group.add_argument('--one-logger-project', type=str, default='e2e-tracking', - help='The one-logger project name. Will ignore if ' - '--enable-one-logger is not set') - group.add_argument('--one-logger-entity', type=str, default='hwinf_dcm', - help='The one-logger username or team name. Will ignore if ' - '--enable-one-logger is not set') - group.add_argument('--one-logger-run-name', type=str, default=None, - help='The one-logger run name displayed. Will ignore if ' - '--enable-one-logger is not set') - return parser - - -def _add_regularization_args(parser): - group = parser.add_argument_group(title='regularization') - - group.add_argument('--attention-dropout', type=float, default=0.1, - help='Post attention dropout probability.') - group.add_argument('--hidden-dropout', type=float, default=0.1, - help='Dropout probability for hidden state transformer.') - group.add_argument('--weight-decay', type=float, default=0.01, - help='Weight decay coefficient for L2 regularization.') - group.add_argument('--start-weight-decay', type=float, - help='Initial weight decay coefficient for L2 regularization.') - group.add_argument('--end-weight-decay', type=float, - help='End of run weight decay coefficient for L2 regularization.') - group.add_argument('--weight-decay-incr-style', type=str, default='constant', - choices=['constant', 'linear', 'cosine'], - help='Weight decay increment function.') - group.add_argument('--clip-grad', type=float, default=1.0, - help='Gradient clipping based on global L2 norm.') - group.add_argument('--adam-beta1', type=float, default=0.9, - help='First coefficient for computing running averages ' - 'of gradient and its square') - group.add_argument('--adam-beta2', type=float, default=0.999, - help='Second coefficient for computing running averages ' - 'of gradient and its square') - group.add_argument('--adam-eps', type=float, default=1e-08, - help='Term added to the denominator to improve' - 'numerical stability') - group.add_argument('--sgd-momentum', type=float, default=0.9, - help='Momentum factor for sgd') - return parser - - -def _add_training_args(parser): - group = parser.add_argument_group(title='training') - - group.add_argument('--micro-batch-size', type=int, default=None, - help='Batch size per model instance (local batch size). ' - 'Global batch size is local batch size times data ' - 'parallel size times number of micro batches.') - group.add_argument('--batch-size', type=int, default=None, - help='Old batch size parameter, do not use. ' - 'Use --micro-batch-size instead') - group.add_argument('--global-batch-size', type=int, default=None, - help='Training batch size. If set, it should be a ' - 'multiple of micro-batch-size times data-parallel-size. ' - 'If this value is None, then ' - 'use micro-batch-size * data-parallel-size as the ' - 'global batch size. This choice will result in 1 for ' - 'number of micro-batches.') - group.add_argument('--rampup-batch-size', nargs='*', default=None, - help='Batch size ramp up with the following values:' - ' --rampup-batch-size ' - ' ' - ' ' - 'For example:' - ' --rampup-batch-size 16 8 300000 \ ' - ' --global-batch-size 1024' - 'will start with global batch size 16 and over ' - ' (1024 - 16) / 8 = 126 intervals will increase' - 'the batch size linearly to 1024. In each interval' - 'we will use approximately 300000 / 126 = 2380 samples.') - group.add_argument('--recompute-activations', action='store_true', - help='recompute activation to allow for training ' - 'with larger models, sequences, and batch sizes.') - group.add_argument('--recompute-granularity', type=str, default=None, - choices=['full', 'selective'], - help='Checkpoint activations to allow for training ' - 'with larger models, sequences, and batch sizes. ' - 'It is supported at two granularities 1) full: ' - 'whole transformer layer is recomputed, ' - '2) selective: core attention part of the transformer ' - 'layer is recomputed.') - group.add_argument('--no-check-for-nan-in-loss-and-grad', action='store_false', - help='Check for NaNs in loss and grad', - dest='check_for_nan_in_loss_and_grad') - group.add_argument('--distribute-saved-activations', - action='store_true', - help='If set, distribute recomputed activations ' - 'across model parallel group.') - group.add_argument('--recompute-method', type=str, default=None, - choices=['uniform', 'block'], - help='1) uniform: uniformly divide the total number of ' - 'Transformer layers and recompute the input activation of ' - 'each divided chunk at specified granularity, ' - '2) recompute the input activations of only a set number of ' - 'individual Transformer layers per pipeline stage and do the ' - 'rest without any recomputing at specified granularity' - 'default) do not apply activations recompute to any layers') - group.add_argument('--recompute-num-layers', type=int, default=None, - help='1) uniform: the number of Transformer layers in each ' - 'uniformly divided recompute unit, ' - '2) block: the number of individual Transformer layers ' - 'to recompute within each pipeline stage.') - group.add_argument('--no-clone-scatter-output-in-embedding', action='store_false', - help='If not set, clone the output of the scatter in embedding layer to GC original tensor.', - dest='clone_scatter_output_in_embedding') - group.add_argument('--profile', action='store_true', - help='Enable nsys profiling. When using this option, nsys ' - 'options should be specified in commandline. An example ' - 'nsys commandline is `nsys profile -s none -t nvtx,cuda ' - '-o --force-overwrite true ' - '--capture-range=cudaProfilerApi ' - '--capture-range-end=stop`.') - group.add_argument('--profile-step-start', type=int, default=10, - help='Global step to start profiling.') - group.add_argument('--profile-step-end', type=int, default=12, - help='Global step to stop profiling.') - group.add_argument('--profile-ranks', nargs='+', type=int, default=[0], - help='Global ranks to profile.') - group.add_argument('--tp-comm-overlap', action='store_true', help='Enables the ' - ' overlap of Tensor parallel communication and GEMM kernels.') - group.add_argument('--tp-comm-overlap-cfg', type=str, default=None, - help='Config file when tp_comm_overlap is enabled.') - group.add_argument('--disable-tp-comm-overlap-ag', action='store_false', - help=('Disables the All-Gather overlap with GEMM by ' - 'pipelining the GEMM and All-Gather.'), - dest='tp_comm_overlap_ag') - group.add_argument('--disable-tp-comm-overlap-rs', action='store_false', - help=('Disables the Reduce-Scatter overlap with GEMM by ' - 'pipelining the GEMM and Reduce-Scatter.'), - dest='tp_comm_overlap_rs') - group.add_argument('--disable-tp-comm-bulk-dgrad', action='store_false', - help='Disables the All-Gather overlap with bprop activation gradient GEMM.', - dest='tp_comm_bulk_dgrad') - group.add_argument('--disable-tp-comm-bulk-wgrad', action='store_false', - help='Disables the Reduce-Scatter overlap with bprop weight gradient GEMM.', - dest='tp_comm_bulk_wgrad') - group.add_argument('--use-cpu-initialization', action='store_true', - default=None, - help='If set, initialize weights on the CPU. This eliminates init differences based on tensor parallelism.') - group.add_argument('--empty-unused-memory-level', default=0, type=int, - choices=[0, 1, 2], - help='Call torch.cuda.empty_cache() each iteration ' - '(training and eval), to reduce fragmentation.' - '0=off, 1=moderate, 2=aggressive.') - - # deprecated - group.add_argument('--checkpoint-activations', action='store_true', - help='Checkpoint activation to allow for training ' - 'with larger models, sequences, and batch sizes.') - group.add_argument('--train-iters', type=int, default=None, - help='Total number of iterations to train over all ' - 'training runs. Note that either train-iters or ' - 'train-samples should be provided.') - group.add_argument('--train-samples', type=int, default=None, - help='Total number of samples to train over all ' - 'training runs. Note that either train-iters or ' - 'train-samples should be provided.') - group.add_argument('--log-interval', type=int, default=100, - help='Report loss and timing interval.') - group.add_argument('--exit-interval', type=int, default=None, - help='Exit the program after the iteration is divisible ' - 'by this value.') - group.add_argument('--exit-duration-in-mins', type=int, default=None, - help='Exit the program after this many minutes.') - group.add_argument('--exit-signal-handler', action='store_true', - help='Dynamically save the checkpoint and shutdown the ' - 'training if SIGTERM is received') - group.add_argument('--tensorboard-dir', type=str, default=None, - help='Write TensorBoard logs to this directory.') - group.add_argument('--no-masked-softmax-fusion', - action='store_false', - help='Disable fusion of query_key_value scaling, ' - 'masking, and softmax.', - dest='masked_softmax_fusion') - group.add_argument('--no-bias-gelu-fusion', action='store_false', - help='Disable bias and gelu fusion.', - dest='bias_gelu_fusion') - group.add_argument('--no-bias-swiglu-fusion', action='store_false', - help='Disable bias and swiglu fusion, the fusion is ' - 'available only when using megatron-core.', - dest='bias_swiglu_fusion') - group.add_argument('--no-bias-dropout-fusion', action='store_false', - help='Disable bias and dropout fusion.', - dest='bias_dropout_fusion') - group.add_argument('--no-rope-fusion', action='store_false', - help='Disable rope fusion, the fusion is available ' - 'only when using megatron-core.', - dest='apply_rope_fusion') - group.add_argument('--use-flash-attn', action='store_true', - help='use FlashAttention implementation of attention. ' - 'https://arxiv.org/abs/2205.14135') - group.add_argument('--disable-bias-linear', action='store_false', - help='Disable bias in the linear layers', - dest='add_bias_linear') - group.add_argument('--add-qkv-bias', action='store_true', - help='Enable bias only in the QKV linear layers', - dest='add_qkv_bias') - group.add_argument('--optimizer', type=str, default='adam', - choices=['adam', 'sgd'], - help='Optimizer function') - group.add_argument('--dataloader-type', type=str, default=None, - choices=['single', 'cyclic', 'external'], - help='Single pass vs multiple pass data loader') - group.add_argument('--no-async-tensor-model-parallel-allreduce', - action='store_false', - help='Disable asynchronous execution of ' - 'tensor-model-parallel all-reduce with weight ' - 'gradient compuation of a column-linear layer.', - dest='async_tensor_model_parallel_allreduce') - group.add_argument('--no-persist-layer-norm', action='store_true', - help='Disable using persistent fused layer norm kernel. ' - 'This kernel supports only a set of hidden sizes. Please ' - 'check persist_ln_hidden_sizes if your hidden ' - 'size is supported.') - group.add_argument('--sequence-parallel', action='store_true', - help='Enable sequence parallel optimization.') - group.add_argument('--no-gradient-accumulation-fusion', - action='store_false', - help='Disable fusing gradient accumulation to weight ' - 'gradient computation of linear layers', - dest='gradient_accumulation_fusion') - group.add_argument('--use-mcore-models', action='store_true', - help='Use the implementation from megatron core') - group.add_argument('--manual-gc', action='store_true', - help='Disable the threshold-based default garbage ' - 'collector and trigger the garbage collection manually. ' - 'Manual garbage collection helps to align the timing of ' - 'the collection across ranks which mitigates the impact ' - 'of CPU-associated jitters. When the manual gc is enabled, ' - 'garbage collection is performed only at the start and the ' - 'end of the validation routine by default.') - group.add_argument('--manual-gc-interval', type=int, default=0, - help='Training step interval to trigger manual garbage ' - 'collection. When the value is set to 0, garbage ' - 'collection is not triggered between training steps.') - group.add_argument('--no-manual-gc-eval', action='store_false', - help='When using manual garbage collection, disable ' - 'garbage collection at the start and the end of each ' - 'evaluation run.', dest='manual_gc_eval') - group.add_argument('--disable-tp-comm-split-ag', action='store_false', - help='Disables the All-Gather overlap with fprop GEMM.', - dest='tp_comm_split_ag') - group.add_argument('--disable-tp-comm-split-rs', action='store_false', - help='Disables the Reduce-Scatter overlap with fprop GEMM.', - dest='tp_comm_split_rs') - - return parser - - -def _add_initialization_args(parser): - group = parser.add_argument_group(title='initialization') - - group.add_argument('--seed', type=int, default=1234, - help='Random seed used for python, numpy, ' - 'pytorch, and cuda.') - group.add_argument('--data-parallel-random-init', action='store_true', - help='Enable random initialization of params ' - 'across data parallel ranks') - group.add_argument('--init-method-std', type=float, default=0.02, - help='Standard deviation of the zero mean normal ' - 'distribution used for weight initialization.') - group.add_argument('--init-method-xavier-uniform', action='store_true', - help='Enable Xavier uniform parameter initialization') - - return parser - - -def _add_learning_rate_args(parser): - group = parser.add_argument_group(title='learning rate') - - group.add_argument('--lr', type=float, default=None, - help='Initial learning rate. Depending on decay style ' - 'and initial warmup, the learning rate at each ' - 'iteration would be different.') - group.add_argument('--lr-decay-style', type=str, default='linear', - choices=['constant', 'linear', 'cosine', 'inverse-square-root'], - help='Learning rate decay function.') - group.add_argument('--lr-decay-iters', type=int, default=None, - help='number of iterations to decay learning rate over,' - ' If None defaults to `--train-iters`') - group.add_argument('--lr-decay-samples', type=int, default=None, - help='number of samples to decay learning rate over,' - ' If None defaults to `--train-samples`') - group.add_argument('--lr-warmup-fraction', type=float, default=None, - help='fraction of lr-warmup-(iters/samples) to use ' - 'for warmup (as a float)') - group.add_argument('--lr-warmup-iters', type=int, default=0, - help='number of iterations to linearly warmup ' - 'learning rate over.') - group.add_argument('--lr-warmup-samples', type=int, default=0, - help='number of samples to linearly warmup ' - 'learning rate over.') - group.add_argument('--lr-warmup-init', type=float, default=0.0, - help='Initial value for learning rate warmup. The ' - 'scheduler starts warmup from this value.') - group.add_argument('--warmup', type=int, default=None, - help='Old lr warmup argument, do not use. Use one of the' - '--lr-warmup-* arguments above') - group.add_argument('--min-lr', type=float, default=0.0, - help='Minimum value for learning rate. The scheduler' - 'clip values below this threshold.') - group.add_argument('--override-opt_param-scheduler', action='store_true', - help='Reset the values of the scheduler (learning rate,' - 'warmup iterations, minimum learning rate, maximum ' - 'number of iterations, and decay style from input ' - 'arguments and ignore values from checkpoints. Note' - 'that all the above values will be reset.') - group.add_argument('--use-checkpoint-opt_param-scheduler', action='store_true', - help='Use checkpoint to set the values of the scheduler ' - '(learning rate, warmup iterations, minimum learning ' - 'rate, maximum number of iterations, and decay style ' - 'from checkpoint and ignore input arguments.') - group.add_argument('--decoupled-lr', type=float, default=None, - help='Separate learning rate for the input and output layer') - group.add_argument('--decoupled-min-lr', type=float, default=None, - help='Minimum value for learning rate for the input and output layer. The scheduler' - 'clip values below this threshold') - - return parser - - -def _add_checkpointing_args(parser): - group = parser.add_argument_group(title='checkpointing') - - group.add_argument('--save', type=str, default=None, - help='Output directory to save checkpoints to.') - group.add_argument('--save-interval', type=int, default=None, - help='Number of iterations between checkpoint saves.') - group.add_argument('--no-save-optim', action='store_true', default=None, - help='Do not save current optimizer.') - group.add_argument('--no-save-rng', action='store_true', default=None, - help='Do not save current rng state.') - group.add_argument('--load', type=str, default=None, - help='Directory containing a model checkpoint.') - group.add_argument('--no-load-optim', action='store_true', default=None, - help='Do not load optimizer when loading checkpoint.') - group.add_argument('--no-load-rng', action='store_true', default=None, - help='Do not load rng state when loading checkpoint.') - group.add_argument('--finetune', action='store_true', - help='Load model for finetuning. Do not load optimizer ' - 'or rng state from checkpoint and set iteration to 0. ' - 'Assumed when loading a release checkpoint.') - group.add_argument('--pretrained-checkpoint', type=str, default=None, - help='Directory containing a pretrained model checkpoint for finetuning.') - group.add_argument('--ckpt-step', type=int, default=None, - help='Checkpoint step to load model from.') - group.add_argument('--no-initialization', action='store_false', - help='Do not perform initialization when building model, ' - 'can reduce startup time when definitely loading from a ' - 'checkpoint', - dest='perform_initialization') - group.add_argument('--use-checkpoint-args', action='store_true', - help='Override any command line arguments with arguments ' - 'from the checkpoint') - group.add_argument('--exit-on-missing-checkpoint', action='store_true', - help="If '--load' is set, but checkpoint is not found " - "(e.g., path typo), then exit instead of random " - "initialization.") - group.add_argument('--use-dist-ckpt', action='store_true', - help='Use distributed checkpoint format.') - group.add_argument('--auto-detect-ckpt-format', action='store_true', - help='Determine if the checkpoint format is in legacy or distributed format.' - ' If False, expects distributed checkpoint iff args.use_dist_ckpt.' - ' Might slow down loading a bit (double rank0 ckpt load).') - group.add_argument('--dist-ckpt-format', type=str, default='torch_dist', - choices=['zarr', 'torch_dist'], - help='Distributed checkpoint format to use.') - group.add_argument('--ckpt-fully-parallel-save', action='store_true', - help='Apply full save parallelization across DP for' - ' distributed checkpoints. Depending on ckpt format' - ' might increase number of files in the checkpoint.') - - return parser - - -def _add_mixed_precision_args(parser): - group = parser.add_argument_group(title='mixed precision') - - group.add_argument('--fp16', action='store_true', - help='Run model in fp16 mode.') - group.add_argument('--bf16', action='store_true', - help='Run model in bfloat16 mode.') - group.add_argument('--loss-scale', type=float, default=None, - help='Static loss scaling, positive power of 2 ' - 'values can improve fp16 convergence. If None, dynamic' - 'loss scaling is used.') - group.add_argument('--initial-loss-scale', type=float, default=2**32, - help='Initial loss-scale for dynamic loss scaling.') - group.add_argument('--min-loss-scale', type=float, default=1.0, - help='Minimum loss scale for dynamic loss scaling.') - group.add_argument('--loss-scale-window', type=float, default=1000, - help='Window over which to raise/lower dynamic scale.') - group.add_argument('--hysteresis', type=int, default=2, - help='hysteresis for dynamic loss scaling') - group.add_argument('--fp32-residual-connection', action='store_true', - help='Move residual connections to fp32.') - group.add_argument('--apply-query-key-layer-scaling', action='store_true', - help='Scale Q * K^T by 1 / layer-number. ' - 'Useful for fp16 training.') - group.add_argument('--attention-softmax-in-fp32', action='store_true', - help='Run attention masking and softmax in fp32. ' - 'This flag is ignored unless ' - '--no-query-key-layer-scaling is specified.') - group.add_argument('--accumulate-allreduce-grads-in-fp32', - action='store_true', - help='Gradient accumulation and all-reduce in fp32.') - group.add_argument('--fp16-lm-cross-entropy', action='store_true', - help='Move the cross entropy unreduced loss calculation' - 'for lm head to fp16.') - - return parser - - -def _add_distributed_args(parser): - group = parser.add_argument_group(title='distributed') - - group.add_argument('--tensor-model-parallel-size', type=int, default=1, - help='Degree of tensor model parallelism.') - group.add_argument('--pipeline-model-parallel-size', type=int, default=1, - help='Degree of pipeline model parallelism.') - group.add_argument('--pipeline-model-parallel-split-rank', - type=int, default=None, - help='Rank where encoder and decoder should be split.') - group.add_argument('--model-parallel-size', type=int, default=None, - help='Old model parallel argument, do not use. Use ' - '--tensor-model-parallel-size instead.') - group.add_argument('--num-layers-per-virtual-pipeline-stage', type=int, default=None, - help='Number of layers per virtual pipeline stage') - group.add_argument('--no-overlap-p2p-communication', action='store_false', - help='overlap pipeline parallel communication with forward and backward chunks', - dest='overlap_p2p_comm') - group.add_argument('--distributed-backend', default='nccl', - choices=['nccl', 'gloo'], - help='Which backend to use for distributed training.') - group.add_argument('--distributed-timeout-minutes', type=int, default=10, - help='Timeout minutes for torch.distributed.') - group.add_argument('--overlap-grad-reduce', action='store_true', - default=False, help='If set, overlap DDP grad reduce.') - group.add_argument('--no-delay-grad-reduce', action='store_false', - help='If not set, delay / synchronize grad reductions in all but first PP stage.', - dest='delay_grad_reduce') - group.add_argument('--overlap-param-gather', action='store_true', - default=False, help='If set, overlap param all-gather in distributed optimizer.') - group.add_argument('--delay-param-gather', action='store_true', - default=False, help='If set, delay / synchronize param all-gathers in all but first PP stage.') - group.add_argument('--no-scatter-gather-tensors-in-pipeline', action='store_false', - help='If not set, use scatter/gather to optimize communication of tensors in pipeline.', - dest='scatter_gather_tensors_in_pipeline') - group.add_argument('--use-ring-exchange-p2p', action='store_true', - default=False, help='If set, use custom-built ring exchange ' - 'for p2p communications. Note that this option will require ' - 'a custom built image that support ring-exchange p2p.') - group.add_argument('--local_rank', type=int, default=None, - help='local rank passed from distributed launcher.') - group.add_argument('--lazy-mpu-init', type=bool, required=False, - help='If set to True, initialize_megatron() ' - 'skips DDP initialization and returns function to ' - 'complete it instead.Also turns on ' - '--use-cpu-initialization flag. This is for ' - 'external DDP manager.' ) - group.add_argument('--standalone-embedding-stage', action='store_true', - default=False, help='If set, *input* embedding layer ' - 'is placed on its own pipeline stage, without any ' - 'transformer layers. (For T5, this flag currently only ' - 'affects the encoder embedding.)') - group.add_argument('--use-distributed-optimizer', action='store_true', - help='Use distributed optimizer.') - group.add_argument('--context-parallel-size', type=int, default=1, - help='Degree of context parallelism.') - group.add_argument('--nccl-communicator-config-path', type=str, default=None, - help='Path to the yaml file with NCCL communicator ' - 'configurations. The number of min/max thread groups and thread ' - 'group cluster size of each communicator can be configured by ' - 'setting `min_ctas`, `max_ctas`, and `cga_cluster_size`.') - return parser - - -def _add_validation_args(parser): - group = parser.add_argument_group(title='validation') - - group.add_argument('--eval-iters', type=int, default=100, - help='Number of iterations to run for evaluation' - 'validation/test for.') - group.add_argument('--eval-interval', type=int, default=1000, - help='Interval between running evaluation on ' - 'validation set.') - group.add_argument("--test-mode", action="store_true", help='Run all real-time test alongside the experiment.') - group.add_argument('--skip-train', action='store_true', - default=False, help='If set, bypass the training loop, ' - 'optionally do evaluation for validation/test, and exit.') - - return parser - - -def _add_data_args(parser): - group = parser.add_argument_group(title='data and dataloader') - - group.add_argument('--data-path', nargs='*', default=None, - help='Path to the training dataset. Accepted format:' - '1) a single data path, 2) multiple datasets in the' - 'form: dataset1-weight dataset1-path dataset2-weight ' - 'dataset2-path ... It is used with --split when a ' - 'single dataset used for all three: train, valid ' - 'and test. It is exclusive to the other ' - '--*-data-path args') - group.add_argument('--split', type=str, default='969, 30, 1', - help='Comma-separated list of proportions for training,' - ' validation, and test split. For example the split ' - '`90,5,5` will use 90%% of data for training, 5%% for ' - 'validation and 5%% for test.') - group.add_argument('--train-data-path', nargs='*', default=None, - help='Path to the training dataset. Accepted format:' - '1) a single data path, 2) multiple datasets in the' - 'form: dataset1-weight dataset1-path dataset2-weight ' - 'dataset2-path ...') - group.add_argument('--valid-data-path', nargs='*', default=None, - help='Path to the validation dataset. Accepted format:' - '1) a single data path, 2) multiple datasets in the' - 'form: dataset1-weight dataset1-path dataset2-weight ' - 'dataset2-path ...') - group.add_argument('--test-data-path', nargs='*', default=None, - help='Path to the test dataset. Accepted format:' - '1) a single data path, 2) multiple datasets in the' - 'form: dataset1-weight dataset1-path dataset2-weight ' - 'dataset2-path ...') - group.add_argument('--data-cache-path', default=None, - help='Path to a directory to hold cached index files.') - group.add_argument('--no-mmap-bin-files', action='store_false', - help='Disable mmap-ing of .bin files.', - dest='mmap_bin_files') - group.add_argument('--mock-data', action='store_true', - help='Skip data loading and validation and opt for artificial ' - 'generation of mock data when an implementation is available.') - - group.add_argument('--vocab-size', type=int, default=None, - help='Size of vocab before EOD or padding.') - group.add_argument('--vocab-file', type=str, default=None, - help='Path to the vocab file.') - group.add_argument('--merge-file', type=str, default=None, - help='Path to the BPE merge file.') - group.add_argument('--vocab-extra-ids', type=int, default=0, - help='Number of additional vocabulary tokens. ' - 'They are used for span masking in the T5 model') - group.add_argument('--seq-length', type=int, default=None, - help='Maximum sequence length to process.') - group.add_argument('--encoder-seq-length', type=int, default=None, - help='Maximum encoder sequence length to process.' - 'This should be exclusive of --seq-length') - group.add_argument('--decoder-seq-length', type=int, default=None, - help="Maximum decoder sequence length to process.") - group.add_argument('--retriever-seq-length', type=int, default=256, - help='Maximum sequence length for the biencoder model ' - 'for retriever') - group.add_argument('--sample-rate', type=float, default=1.0, - help='sample rate for training data. Supposed to be 0 ' - ' < sample_rate < 1') - group.add_argument('--mask-prob', type=float, default=0.15, - help='Probability of replacing a token with mask.') - group.add_argument('--short-seq-prob', type=float, default=0.1, - help='Probability of producing a short sequence.') - group.add_argument('--num-workers', type=int, default=2, - help="Dataloader number of workers.") - group.add_argument('--tokenizer-type', type=str, - default=None, - choices=['BertWordPieceLowerCase', - 'BertWordPieceCase', - 'GPT2BPETokenizer', - 'SentencePieceTokenizer', - 'GPTSentencePieceTokenizer', - 'Llama2Tokenizer', - 'Llama3Tokenizer', - 'NullTokenizer'], - help='What type of tokenizer to use.') - group.add_argument('--tokenizer-model', type=str, default=None, - help='Sentencepiece tokenizer model.') - group.add_argument('--reset-position-ids', action='store_true', - help='Reset posistion ids after end-of-document token.') - group.add_argument('--reset-attention-mask', action='store_true', - help='Reset self attention maske after ' - 'end-of-document token.') - group.add_argument('--eod-mask-loss', action='store_true', - help='Mask loss for the end of document tokens.') - group.add_argument('--no-create-attention-mask-in-dataloader', action='store_false', - help='If set, do not create attention_masks in dataloader.', - dest='create_attention_mask_in_dataloader') - - return parser - - -def _add_autoresume_args(parser): - group = parser.add_argument_group(title='autoresume') - - group.add_argument('--adlr-autoresume', action='store_true', - help='Enable autoresume on adlr cluster.') - group.add_argument('--adlr-autoresume-interval', type=int, default=1000, - help='Intervals over which check for autoresume' - 'termination signal') - - return parser - - -def _add_biencoder_args(parser): - group = parser.add_argument_group(title='biencoder') - - # network size - group.add_argument('--ict-head-size', type=int, default=None, - help='Size of block embeddings to be used in ICT and ' - 'REALM (paper default: 128)') - group.add_argument('--biencoder-projection-dim', type=int, default=0, - help='Size of projection head used in biencoder (paper' - ' default: 128)') - group.add_argument('--biencoder-shared-query-context-model', action='store_true', - help='Whether to share the parameters of the query ' - 'and context models or not') - - # checkpointing - group.add_argument('--ict-load', type=str, default=None, - help='Directory containing an ICTBertModel checkpoint') - group.add_argument('--bert-load', type=str, default=None, - help='Directory containing an BertModel checkpoint ' - '(needed to start ICT and REALM)') - - # data - group.add_argument('--titles-data-path', type=str, default=None, - help='Path to titles dataset used for ICT') - group.add_argument('--query-in-block-prob', type=float, default=0.1, - help='Probability of keeping query in block for ' - 'ICT dataset') - group.add_argument('--use-one-sent-docs', action='store_true', - help='Whether to use one sentence documents in ICT') - group.add_argument('--evidence-data-path', type=str, default=None, - help='Path to Wikipedia Evidence frm DPR paper') - - # training - group.add_argument('--retriever-report-topk-accuracies', nargs='+', type=int, - default=[], help="Which top-k accuracies to report " - "(e.g. '1 5 20')") - group.add_argument('--retriever-score-scaling', action='store_true', - help='Whether to scale retriever scores by inverse ' - 'square root of hidden size') - - # faiss index - group.add_argument('--block-data-path', type=str, default=None, - help='Where to save/load BlockData to/from') - group.add_argument('--embedding-path', type=str, default=None, - help='Where to save/load Open-Retrieval Embedding' - ' data to/from') - - # indexer - group.add_argument('--indexer-batch-size', type=int, default=128, - help='How large of batches to use when doing indexing ' - 'jobs') - group.add_argument('--indexer-log-interval', type=int, default=1000, - help='After how many batches should the indexer ' - 'report progress') - return parser - - -def _add_vision_args(parser): - group = parser.add_argument_group(title="vision") - - # general vision arguements - group.add_argument('--num-classes', type=int, default=1000, - help='num of classes in vision classificaiton task') - group.add_argument('--img-h', type=int, default=224, - help='Image height for vision classification task') - group.add_argument('--img-w', type=int, default=224, - help='Image height for vision classification task') - group.add_argument('--num-channels', type=int, default=3, - help='Number of channels in input image data') - group.add_argument('--patch-dim', type=int, default=16, - help='patch dimension') - group.add_argument('--classes-fraction', type=float, default=1.0, - help='training with fraction of classes.') - group.add_argument('--data-per-class-fraction', type=float, default=1.0, - help='training with fraction of data per class.') - group.add_argument('--no-data-sharding', action='store_false', - help='Disable data sharding.', - dest='data_sharding') - group.add_argument('--head-lr-mult', type=float, default=1.0, - help='learning rate multiplier for head during finetuning') - - # pretraining type and backbone selection` - group.add_argument('--vision-pretraining', action='store_true', - help='flag to indicate vision pretraining') - group.add_argument('--vision-pretraining-type', type=str, default='classify', - choices=['classify', 'inpaint', 'dino'], - help='pretraining objectives') - group.add_argument('--vision-backbone-type', type=str, default='vit', - choices=['vit', 'mit', 'swin'], - help='backbone types types') - group.add_argument('--swin-backbone-type', type=str, default='tiny', - choices=['tiny', 'base', 'h3'], - help='pretraining objectives') - # inpainting arguments - group.add_argument('--mask-type', type=str, default='random', - choices=['random', 'row'], - help='mask types') - group.add_argument('--mask-factor', type=float, default=1.0, - help='mask size scaling parameter') - - # dino arguments - group.add_argument('--iter-per-epoch', type=int, default=1250, - help='iterations per epoch') - group.add_argument('--dino-local-img-size', type=int, default=96, - help='Image size for vision classification task') - group.add_argument('--dino-local-crops-number', type=int, default=10, - help='Number of local crops') - group.add_argument('--dino-head-hidden-size', type=int, default=2048, - help='Hidden dimension size in dino head') - group.add_argument('--dino-bottleneck-size', type=int, default=256, - help='Bottle neck dimension in dino head ') - group.add_argument('--dino-freeze-last-layer', type=float, default=1, - help='Freezing last layer weights') - group.add_argument('--dino-norm-last-layer', action='store_true', - help='Disable Norm in last layer.') - group.add_argument('--dino-warmup-teacher-temp', type=float, default=0.04, - help='warump teacher temperature') - group.add_argument('--dino-teacher-temp', type=float, default=0.07, - help='teacher temperature') - group.add_argument('--dino-warmup-teacher-temp-epochs', type=int, default=30, - help='warmup teacher temperaure epochs') - - # regularization arguments - group.add_argument('--qk-layernorm', action='store_true', - help='Whether to layer normalize the q and k attention embeddings.') - - return parser - -def _add_moe_args(parser): - group = parser.add_argument_group(title="moe") - group.add_argument('--expert-model-parallel-size', type=int, default=1, - help='Degree of expert model parallelism.') - group.add_argument('--num-experts', type=int, default=None, - help='Number of Experts in MoE (None means no MoE)') - group.add_argument('--moe-router-load-balancing-type', type=str, - choices=['aux_loss', 'sinkhorn', "none"], - default='aux_loss', - help='Determines the load balancing strategy for the router. "aux_loss" corresponds to the load balancing loss used in GShard and SwitchTransformer, "sinkhorn" corresponds to the balancing algorithm used in S-BASE, and "none" implies no load balancing. The default is "aux_loss".') - group.add_argument('--moe-router-topk', type=int, default=2, - help='Number of experts to route to for each token. The default is 2.') - group.add_argument('--moe-grouped-gemm', action='store_true', - help='When there are multiple experts per rank, compress multiple local (potentially small) gemms in a single kernel launch to improve the utilization and performance by leveraging the Grouped GEMM feature introduced since CUTLASS 2.8 (https://github.com/fanshiqing/grouped_gemm).') - group.add_argument('--moe-aux-loss-coeff', type=float, default=0.0, - help='Scaling coefficient for the aux loss: a starting value of 1e-2 is recommended.') - group.add_argument('--moe-z-loss-coeff', type=float, default=None, - help='Scaling coefficient for the z-loss: a starting value of 1e-3 is recommended.') - group.add_argument('--moe-input-jitter-eps', type=float, default=None, - help='Add noise to the input tensor by applying jitter with a specified epsilon value.') - group.add_argument('--moe-token-dropping', action='store_true', - help='This feature involves selectively dropping and padding tokens for each expert to achieve a specified capacity, similar to GShard, Switch-Transformer, and DeepSpeed-MoE. Note: Currently unsupported.') - group.add_argument('--moe-token-dispatcher-type', type=str, - choices=['allgather', 'alltoall'], - default='allgather', - help='.') - group.add_argument('--moe-per-layer-logging', action='store_true', - help='Enable per-layer logging for MoE, currently supports auxiliary loss and z loss.') - - return parser - -def _add_experimental_args(parser): - group = parser.add_argument_group(title='experimental') - - group.add_argument('--spec', type=str, default=None, nargs='*', - help='Specify the pair ' - 'that returns a spec to customize a model, transformer ' - 'block, or transformer layer, depending on the use case.' - 'To use local spec specify local as the argument.' - 'For more details, see the model class, ' - '`transformer_block.py`, or `transformer_layer.py`') - group.add_argument('--yaml-cfg', type=str, default=None, - help = 'Config file to add additional arguments') - - return parser diff --git a/training/benchmarks/llama3_8B/megatron/pretrain_llama3.sh b/training/benchmarks/llama3_8B/megatron/pretrain_llama3.sh index ee1eca7ff..4eedec946 100755 --- a/training/benchmarks/llama3_8B/megatron/pretrain_llama3.sh +++ b/training/benchmarks/llama3_8B/megatron/pretrain_llama3.sh @@ -1,10 +1,29 @@ #!/bin/bash +# Copyright (c) 2024 BAAI. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License") +#!/usr/bin/env python3 +# -*- coding: UTF-8 -*- +# args + +DATA_DIR=$1 +GPUS_PER_NODE=$2 +NNODES=$3 +NODE_RANK=$4 +MASTER_ADDR=$5 +MASTER_PORT=$6 +MEGAPATH=$7 +MBS=$8 +ITERS=$9 +TP=${10} +PP=${11} +TDIR=${12} +ADAPT=${13} # Runs the "LLaMA3-8B" parameter model -export PYTHONPATH=$PYTHONPATH:/workspace/Megatron-LM +export PYTHONPATH=$PYTHONPATH:$MEGAPATH export CUDA_DEVICE_MAX_CONNECTIONS=1 -ls /workspace # algorithm args MODEL_ARGS=" \ @@ -24,8 +43,8 @@ MODEL_ARGS=" \ OPT_ARGS=" \ --lr 1.0e-5 \ - --train-iters 300 \ - --lr-decay-iters 300 \ + --train-iters $ITERS \ + --lr-decay-iters $ITERS \ --lr-decay-style cosine \ --min-lr 1.0e-6 \ --weight-decay 0.1 \ @@ -39,17 +58,26 @@ ALGO_ARGS="$MODEL_ARGS $OPT_ARGS" # data args DATA_ARGS=" \ - --data-path /data/llama3_8b_pretrain/wudao_llama3bpe_content_document \ + --data-path $DATA_DIR/wudao_llama3bpe_content_document \ --tokenizer-type Llama3Tokenizer \ + --tokenizer-model $TDIR \ --split 100,0,0 " # training args +DISTRIBUTED_ARGS=" + --nproc_per_node $GPUS_PER_NODE \ + --nnodes $NNODES \ + --node_rank $NODE_RANK \ + --master_addr $MASTER_ADDR \ + --master_port $MASTER_PORT +" + TRAINING_ARGS=" \ - --micro-batch-size 1 \ - --tensor-model-parallel-size 1 \ - --pipeline-model-parallel-size 2 \ + --micro-batch-size $MBS \ + --tensor-model-parallel-size $TP \ + --pipeline-model-parallel-size $PP \ --sequence-parallel \ --bf16 " @@ -65,7 +93,8 @@ VENDOR_ARGS=" \ OUTPUT_ARGS=" --log-interval 1" -run_cmd="torchrun --nproc_per_node 8 --nnodes 1 --node_rank 0 --master_addr 10.1.2.59 --master_port 29501 /workspace/Megatron-LM/pretrain_gpt.py \ +source $ADAPT +run_cmd="torchrun $DISTRIBUTED_ARGS $MEGAPATH/pretrain_gpt.py \ $ALGO_ARGS \ $DATA_ARGS \ $TRAINING_ARGS \ diff --git a/training/benchmarks/llama3_8B/megatron/run_pretraining.py b/training/benchmarks/llama3_8B/megatron/run_pretraining.py index e076054fe..620c77d82 100644 --- a/training/benchmarks/llama3_8B/megatron/run_pretraining.py +++ b/training/benchmarks/llama3_8B/megatron/run_pretraining.py @@ -1,3 +1,8 @@ +# Copyright (c) 2024 BAAI. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License") +#!/usr/bin/env python3 +# -*- coding: UTF-8 -*- import subprocess from argparse import ArgumentParser import os @@ -35,19 +40,56 @@ def parse_args(): module = import_module(config_file) - seqlength = getattr(module, 'seqlength') localbs = getattr(module, 'localbs') train_steps = getattr(module, 'train_steps') theoryflops = getattr(module, 'theoryflops') - epochs = getattr(module, 'epochs') + megapath = getattr(module, 'megatron_path') tensor_parallel = getattr(module, 'tensor_parallel') pipeline_parallel = getattr(module, 'pipeline_parallel') + tokenizer = getattr(module, 'tokenizer_path') + + nnodes = args.nnodes + nproc_per_node = args.nproc_per_node + node_rank = args.node_rank + master_addr = args.master_addr + master_port = args.master_port + data_dir = args.data_dir + tokenizer_dir = os.path.join(args.data_dir, tokenizer) task_log_file = os.path.join(args.log_dir, "megatron.log.txt") + + # merge llama3 patch - exec_cmd = "cp arguments.py /workspace/Megatron-LM/megatron/training/;cp tokenizer.py /workspace/Megatron-LM/megatron/training/tokenizer/;bash pretrain_llama3.sh" + + origin_file = os.path.join(megapath, "megatron/training/arguments.py") + exec_cmd = "patch --silent --forward " + origin_file + " arguments.patch -o tmp.py;mv tmp.py " + origin_file + exec_cmd = exec_cmd + ";" + + origin_file = os.path.join(megapath, "megatron/training/tokenizer/tokenizer.py") + exec_cmd = exec_cmd + "patch --silent --forward " + origin_file + " tokenizer.patch -o tmp.py;mv tmp.py " + origin_file + exec_cmd = exec_cmd + ";" + + # bash pretrain_llama3.sh + + exec_cmd = exec_cmd + "bash pretrain_llama3.sh" + + # args + + exec_cmd = exec_cmd + " " + data_dir + exec_cmd = exec_cmd + " " + str(nproc_per_node) + exec_cmd = exec_cmd + " " + str(nnodes) + exec_cmd = exec_cmd + " " + str(node_rank) + exec_cmd = exec_cmd + " " + str(master_addr) + exec_cmd = exec_cmd + " " + str(master_port) + exec_cmd = exec_cmd + " " + megapath + exec_cmd = exec_cmd + " " + str(localbs) + exec_cmd = exec_cmd + " " + str(train_steps) + exec_cmd = exec_cmd + " " + str(tensor_parallel) + exec_cmd = exec_cmd + " " + str(pipeline_parallel) + exec_cmd = exec_cmd + " " + str(tokenizer_dir) exec_cmd = exec_cmd + " " + os.path.join(config_dir_path, "training_adapter.sh") + print(exec_cmd) with open(task_log_file, "w") as f: p = subprocess.Popen(exec_cmd, @@ -64,8 +106,8 @@ def parse_args(): steptime = info.split(":")[1] time_per_step = float(steptime) / 1000 - whole_tps = 512* seqlength / time_per_step + whole_tps = 512 * 8192 / time_per_step chip_tps = whole_tps / (args.nproc_per_node * args.nnodes) print("System tokens per second: ", whole_tps) print("Tokens/p/s: ", chip_tps) - print("MFU: ", chip_tps * 80000000000.0 * 6 / theoryflops) + print("MFU: ", chip_tps * 8000000000.0 * 6 / theoryflops) diff --git a/training/benchmarks/llama3_8B/megatron/tokenizer.patch b/training/benchmarks/llama3_8B/megatron/tokenizer.patch new file mode 100644 index 000000000..3c131e2cf --- /dev/null +++ b/training/benchmarks/llama3_8B/megatron/tokenizer.patch @@ -0,0 +1,52 @@ +# Copyright (c) 2024 BAAI. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License") +#!/usr/bin/env python3 +# -*- coding: UTF-8 -*- +43a44,45 +> elif args.tokenizer_type == 'Llama3Tokenizer': +> tokenizer = _Llama3Tokenizer(args.tokenizer_model) +439a442,484 +> +> class _Llama3Tokenizer(MegatronTokenizer): +> def __init__(self, model_file): +> super().__init__("placeholder") +> from transformers import AutoTokenizer +> self.tokenizer = AutoTokenizer.from_pretrained( +> model_file, +> padding_side="right", +> use_fast=False, +> trust_remote_code=True +> ) +> self.extra_vocab_size = 256 +> +> @property +> def vocab_size(self): +> return self.tokenizer.vocab_size + self.extra_vocab_size +> +> @property +> def vocab(self): +> return self.tokenizer.encoder +> +> @property +> def inv_vocab(self): +> return self.tokenizer.decoder +> +> def tokenize(self, text): +> return self.tokenizer.encode(text) +> +> def detokenize(self, token_ids): +> return self.tokenizer.decode(token_ids) +> +> @property +> def eod(self): +> return self.tokenizer.eos_token_id +> +> @property +> def eos_token(self): +> return self.tokenizer.eos_token +> +> @property +> def pad_token_id(self): +> return self.tokenizer.pad_token_id +> diff --git a/training/benchmarks/llama3_8B/megatron/tokenizer.py b/training/benchmarks/llama3_8B/megatron/tokenizer.py deleted file mode 100644 index b824f84ea..000000000 --- a/training/benchmarks/llama3_8B/megatron/tokenizer.py +++ /dev/null @@ -1,567 +0,0 @@ -# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. - -"""Megatron tokenizers.""" - -from abc import ABC -from abc import abstractmethod - -from megatron.core.datasets.megatron_tokenizer import MegatronTokenizer - -from .bert_tokenization import FullTokenizer as FullBertTokenizer -from .gpt2_tokenization import GPT2Tokenizer - - -def build_tokenizer(args): - """Initialize tokenizer.""" - if args.rank == 0: - print('> building {} tokenizer ...'.format(args.tokenizer_type), - flush=True) - - # Select and instantiate the tokenizer. - if args.tokenizer_type == 'BertWordPieceLowerCase': - assert args.vocab_file is not None - tokenizer = _BertWordPieceTokenizer(vocab_file=args.vocab_file, - lower_case=True, - vocab_extra_ids=args.vocab_extra_ids) - elif args.tokenizer_type == 'BertWordPieceCase': - assert args.vocab_file is not None - tokenizer = _BertWordPieceTokenizer(vocab_file=args.vocab_file, - lower_case=False, - vocab_extra_ids=args.vocab_extra_ids) - elif args.tokenizer_type == 'GPT2BPETokenizer': - assert args.vocab_file is not None - assert args.merge_file is not None - tokenizer = _GPT2BPETokenizer(args.vocab_file, args.merge_file) - elif args.tokenizer_type == 'SentencePieceTokenizer': - assert args.tokenizer_model is not None - tokenizer = _SentencePieceTokenizer(args.tokenizer_model, vocab_extra_ids=args.vocab_extra_ids) - elif args.tokenizer_type == 'GPTSentencePieceTokenizer': - assert args.tokenizer_model is not None - tokenizer = _GPTSentencePieceTokenizer(args.tokenizer_model) - elif args.tokenizer_type == 'Llama2Tokenizer': - assert args.tokenizer_model is not None - tokenizer = _Llama2Tokenizer(args.tokenizer_model) - elif args.tokenizer_type == 'Llama3Tokenizer': - tokenizer = _Llama3Tokenizer() - elif args.tokenizer_type == 'NullTokenizer': - assert args.vocab_size is not None - tokenizer = _NullTokenizer(args.vocab_size) - else: - raise NotImplementedError('{} tokenizer is not ' - 'implemented.'.format(args.tokenizer_type)) - - # Add vocab size (if not already set from a checkpoint). - if getattr(args, "padded_vocab_size", None) is None: - args.padded_vocab_size = _vocab_size_with_padding(tokenizer.vocab_size, - args) - - return tokenizer - - -def _vocab_size_with_padding(orig_vocab_size, args): - """Pad vocab size so it is divisible by model parallel size and - still having GPU friendly size.""" - - after = orig_vocab_size - multiple = args.make_vocab_size_divisible_by * \ - args.tensor_model_parallel_size - while (after % multiple) != 0: - after += 1 - if args.rank == 0: - print(' > padded vocab (size: {}) with {} dummy tokens ' - '(new size: {})'.format( - orig_vocab_size, after - orig_vocab_size, after), flush=True) - return after - - -class _BertWordPieceTokenizer(MegatronTokenizer): - """Original BERT wordpiece tokenizer.""" - - def __init__(self, vocab_file, lower_case=True, vocab_extra_ids=0): - super().__init__(vocab_file, lower_case=lower_case, vocab_extra_ids=vocab_extra_ids) - self.tokenizer = FullBertTokenizer(vocab_file, do_lower_case=lower_case) - self.cls_id = self.tokenizer.vocab['[CLS]'] - self.sep_id = self.tokenizer.vocab['[SEP]'] - self.pad_id = self.tokenizer.vocab['[PAD]'] - self.mask_id = self.tokenizer.vocab['[MASK]'] - self._additional_special_tokens = [] - - # (dsachan) Add BOS and EOS tokens - SPECIAL_TOKENS = {'eos_token': '[EOS]', - 'bos_token': '[BOS]'} - self._bos_token = '[BOS]' - self.add_token(self._bos_token) - self._bos_token_id = self.vocab.get(self._bos_token) - - self._eos_token = '[EOS]' - self.add_token(self._eos_token) - self._eos_token_id = self.vocab.get(self._eos_token) - - # (dsachan) Add additional special tokens - # These can be used as sentinel tokens in T5 model inputs - additional_special_tokens = [] - additional_special_tokens.extend( - ["".format(i) for i in range(vocab_extra_ids)]) - self.add_additional_special_tokens(additional_special_tokens) - - def add_token(self, token): - if token not in self.vocab: - self.inv_vocab[self.vocab_size] = token - # self.vocab_size comes from len(vocab) - # and it will increase as we add elements - self.vocab[token] = self.vocab_size - - def add_additional_special_tokens(self, tokens_list): - setattr(self, "additional_special_tokens", tokens_list) - for value in tokens_list: - self.add_token(value) - - @property - def vocab_size(self): - return self.tokenizer.vocab_size() - - @property - def vocab(self): - return self.tokenizer.vocab - - @property - def inv_vocab(self): - return self.tokenizer.inv_vocab - - def tokenize(self, text): - text_tokens = self.tokenizer.tokenize(text) - return self.tokenizer.convert_tokens_to_ids(text_tokens) - - def decode(self, ids): - tokens = self.tokenizer.convert_ids_to_tokens(ids) - return self.tokenizer.convert_tokens_to_string(tokens) - - def decode_token_ids(self, token_ids): - tokens = self.tokenizer.convert_ids_to_tokens(token_ids) - exclude_list = ['[PAD]', '[CLS]'] - non_pads = [t for t in tokens if t not in exclude_list] - - result = "" - for s in non_pads: - if s.startswith("##"): - result += s[2:] - else: - result += " " + s - - return result - - @property - def cls(self): - return self.cls_id - - @property - def sep(self): - return self.sep_id - - @property - def pad(self): - return self.pad_id - - @property - def mask(self): - return self.mask_id - - @property - def bos(self): - """ Id of the beginning of sentence token in the vocabulary.""" - return self._bos_token_id - - @property - def eos(self): - """ Id of the end of sentence token in the vocabulary.""" - return self._eos_token_id - - @property - def bos_token(self): - """ Beginning of sentence token id """ - return self._bos_token - - @property - def eos_token(self): - """ End of sentence token id """ - return self._eos_token - - @property - def additional_special_tokens(self): - """ All the additional special tokens you may want to use (list of strings).""" - return self._additional_special_tokens - - @property - def additional_special_tokens_ids(self): - """ Ids of all the additional special tokens in the vocabulary (list of integers).""" - return [self.vocab.get(token) for token in self._additional_special_tokens] - - @additional_special_tokens.setter - def additional_special_tokens(self, value): - self._additional_special_tokens = value - - -class _GPT2BPETokenizer(MegatronTokenizer): - """Original GPT2 BPE tokenizer.""" - - def __init__(self, vocab_file, merge_file): - super().__init__(vocab_file, merge_file) - - self.tokenizer = GPT2Tokenizer(vocab_file, merge_file, errors='replace', - special_tokens=[], max_len=None) - self.eod_id = self.tokenizer.encoder['<|endoftext|>'] - - @property - def vocab_size(self): - return len(self.tokenizer.encoder) - - @property - def vocab(self): - return self.tokenizer.encoder - - @property - def inv_vocab(self): - return self.tokenizer.decoder - - def tokenize(self, text): - return self.tokenizer.encode(text) - - def detokenize(self, token_ids): - return self.tokenizer.decode(token_ids) - - @property - def eod(self): - return self.eod_id - - -class _SentencePieceTokenizer(MegatronTokenizer): - """SentencePieceTokenizer-Megatron wrapper""" - - def __init__(self, model_file, vocab_extra_ids=0): - super().__init__(model_file, vocab_extra_ids=vocab_extra_ids) - - import sentencepiece - self.tokenizer = sentencepiece.SentencePieceProcessor(model_file=model_file) - self._initalize(vocab_extra_ids) - - def _populate_vocab(self): - self._vocab = {} - self._inv_vocab = {} - - for i in range(len(self.tokenizer)): - t = self.tokenizer.id_to_piece(i) - self._inv_vocab[i] = t - self._vocab[t] = i - - def _initalize(self, vocab_extra_ids): - self._populate_vocab() - self._special_tokens = {} - self._inv_special_tokens = {} - - self._t5_tokens = [] - - def _add_special_token(t): - if t not in self._vocab: - next_id = len(self._vocab) - self._vocab[t] = next_id - self._inv_vocab[next_id] = t - self._special_tokens[t] = self._vocab[t] - self._inv_special_tokens[self._vocab[t]] = t - - _add_special_token('') - self._cls_id = self._vocab[''] - _add_special_token('') - self._sep_id = self._vocab[''] - _add_special_token('') - self._eod_id = self._vocab[''] - _add_special_token('') - self._mask_id = self._vocab[''] - - pad_id = self.tokenizer.pad_id() - try: - pad_token = self.tokenizer.id_to_piece(pad_id) - except IndexError: - pad_token = '' - _add_special_token(pad_token) - self._pad_id = self._vocab[pad_token] - - bos_id = self.tokenizer.bos_id() - try: - bos_token = self.tokenizer.id_to_piece(bos_id) - except IndexError: - bos_token = '' - _add_special_token(bos_token) - self._bos_id = self._vocab[bos_token] - - eos_id = self.tokenizer.eos_id() - try: - eos_token = self.tokenizer.id_to_piece(eos_id) - except IndexError: - eos_token = '' - _add_special_token(eos_token) - self._eos_id = self._vocab[eos_token] - - for i in range(vocab_extra_ids): - t = "".format(i) - _add_special_token(t) - self._t5_tokens += [t] - - @property - def vocab_size(self): - return len(self._vocab) - - @property - def vocab(self): - return self._vocab - - @property - def inv_vocab(self): - return self._inv_vocab - - @property - def decoder(self): - return self._inv_vocab - - @property - def encoder(self): - return self._vocab - - # From: - # https://github.com/NVIDIA/NeMo/blob/c8fa217e811d60d11d014827c7f3845ff6c99ae7/nemo/collections/common/tokenizers/sentencepiece_tokenizer.py#L89 - def tokenize(self, text): - ids = [] - idx = 0 - - while 1: - indices = {} - for token in self._special_tokens: - try: - indices[token] = text[idx:].index(token) - except ValueError: - continue - if len(indices) == 0: - break - - next_token = min(indices, key=indices.get) - next_idx = idx + indices[next_token] - - ids.extend(self.tokenizer.encode_as_ids(text[idx:next_idx])) - ids.append(self._special_tokens[next_token]) - idx = next_idx + len(next_token) - - ids.extend(self.tokenizer.encode_as_ids(text[idx:])) - return ids - - # From: - # https://github.com/NVIDIA/NeMo/blob/c8fa217e811d60d11d014827c7f3845ff6c99ae7/nemo/collections/common/tokenizers/sentencepiece_tokenizer.py#L125 - def detokenize(self, ids): - text = "" - last_i = 0 - - for i, id in enumerate(ids): - if id in self._inv_special_tokens: - text += self.tokenizer.decode_ids(ids[last_i:i]) + " " - text += self._inv_special_tokens[id] + " " - last_i = i + 1 - - text += self.tokenizer.decode_ids(ids[last_i:]) - return text - - @property - def cls(self): - return self._cls_id - - @property - def sep(self): - return self._sep_id - - @property - def pad(self): - return self._pad_id - - @property - def bos(self): - return self._bos_id - - @property - def eod(self): - return self._eod_id - - @property - def eos(self): - return self._eos_id - - @property - def mask(self): - return self._mask_id - - @property - def additional_special_tokens_ids(self): - return [self.vocab[k] for k in self._t5_tokens] - - -class _GPTSentencePieceTokenizer(_SentencePieceTokenizer): - """SentencePieceTokenizer-Megatron wrapper""" - - def __init__(self, model_file,): - super().__init__(model_file, vocab_extra_ids=0) - - def _initalize(self, vocab_extra_ids): - self._populate_vocab() - - self._pad_id = self.tokenizer.pad_id() - self._bos_id = self.tokenizer.bos_id() - self._eos_id = self.tokenizer.eos_id() - - def tokenize(self, text): - return self.tokenizer.encode_as_ids(text) - - def detokenize(self, ids): - return self.tokenizer.decode_ids(ids) - - @property - def cls(self): - return -1 - - @property - def sep(self): - return -1 - - @property - def mask(self): - return -1 - - @property - def eod(self): - return self._eos_id - - @property - def additional_special_tokens_ids(self): - return None - - -class _Llama3Tokenizer(MegatronTokenizer): - def __init__(self,): - super().__init__("placeholder") - from transformers import AutoTokenizer - self.tokenizer = AutoTokenizer.from_pretrained( - "/data/llama3_8b_pretrain/llama3_8b_hf", - padding_side="right", - use_fast=False, - trust_remote_code=True - ) - self.extra_vocab_size = 256 - - @property - def vocab_size(self): - return self.tokenizer.vocab_size + self.extra_vocab_size - - @property - def vocab(self): - return self.tokenizer.encoder - - @property - def inv_vocab(self): - return self.tokenizer.decoder - - def tokenize(self, text): - return self.tokenizer.encode(text) - - def detokenize(self, token_ids): - return self.tokenizer.decode(token_ids) - - @property - def eod(self): - return self.tokenizer.eos_token_id - - @property - def eos_token(self): - return self.tokenizer.eos_token - - @property - def pad_token_id(self): - return self.tokenizer.pad_token_id - - -class _Llama2Tokenizer(_SentencePieceTokenizer): - """SentencePieceTokenizer-Megatron wrapper""" - - def __init__(self, model_file,): - super().__init__(model_file, vocab_extra_ids=0) - - def _initalize(self, vocab_extra_ids): - self._populate_vocab() - - # BOS / EOS token IDs - self.n_words: int = self.tokenizer.vocab_size() - self.bos_id: int = self.tokenizer.bos_id() - self.eos_id: int = self.tokenizer.eos_id() - self.pad_id: int = self.tokenizer.pad_id() - assert self.tokenizer.vocab_size() == self.tokenizer.get_piece_size() - - def tokenize(self, s: str, bos=True, eos=False): - '''Default args for text completion, not chat/dialog.''' - assert type(s) is str - t = self.tokenizer.encode(s) - if bos: - t = [self.bos_id] + t - if eos: - t = t + [self.eos_id] - return t - - def detokenize(self, ids): - return self.tokenizer.decode_ids(ids) - - @property - def cls(self): - return -1 - - @property - def sep(self): - return -1 - - @property - def mask(self): - return -1 - - @property - def eod(self): - return self.eos_id - - @property - def additional_special_tokens_ids(self): - return None - - -class _NullTokenizer: - def __init__(self, vocab_size): - vocab_size = int(vocab_size) - self._eos_id = vocab_size - self.vocab_size = vocab_size+1 - - def tokenize(self, text): - return [int(x) for x in text.split(' ')] - - def detokenize(self, ids): - text = [str(x) for x in ids] - return ' '.join(text) - - @property - def cls(self): - return -1 - - @property - def sep(self): - return -1 - - @property - def mask(self): - return -1 - - @property - def eod(self): - return self._eos_id - - @property - def additional_special_tokens_ids(self): - return None diff --git a/training/nvidia/docker_image/megatron_core060/megatron_core060_install.sh b/training/nvidia/docker_image/megatron_core060/megatron_core060_install.sh index efca1a78d..a9bf588e2 100644 --- a/training/nvidia/docker_image/megatron_core060/megatron_core060_install.sh +++ b/training/nvidia/docker_image/megatron_core060/megatron_core060_install.sh @@ -1,3 +1 @@ #!/bin/bash -# using github mirrors to avoid github TTL - diff --git a/training/nvidia/llama3_8B-megatron/config/config_A100x1x8.py b/training/nvidia/llama3_8B-megatron/config/config_A100x1x8.py index 9747f2131..71ca6d08a 100644 --- a/training/nvidia/llama3_8B-megatron/config/config_A100x1x8.py +++ b/training/nvidia/llama3_8B-megatron/config/config_A100x1x8.py @@ -1,7 +1,7 @@ -seqlength = 8192 +tokenizer_path = "llama3_8b_hf" localbs = 1 -train_steps = 10 +train_steps = 300 theoryflops = 312000000000000.0 -epochs = 1 -tensor_parallel = 2 -pipeline_parallel = 1 +megatron_path = "/workspace/Megatron-LM" # need to be aligned with DockerFile. In NGCtorch, it's /workspace/ + Megatron-LM +tensor_parallel = 1 +pipeline_parallel = 2 From 11c0a300fd35ece43878e4cad8cb47f1d5b8f1c8 Mon Sep 17 00:00:00 2001 From: shh2000 <13820618441@163.com> Date: Sun, 28 Apr 2024 11:40:30 +0800 Subject: [PATCH 3/7] =?UTF-8?q?add=20lic=20=E2=80=9D?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- base/benchmarks/computation-BF16/main.py | 5 +++++ base/benchmarks/computation-FP16/main.py | 5 +++++ base/benchmarks/computation-FP32/main.py | 5 +++++ base/benchmarks/computation-FP64/main.py | 5 +++++ base/benchmarks/computation-TF32/main.py | 5 +++++ base/benchmarks/drivers/utils.py | 7 ++++++- base/benchmarks/main_memory-bandwidth/main.py | 5 +++++ base/container_main.py | 5 +++++ base/toolkits/computation-BF16/nvidia/gemm.cu | 4 ++++ base/toolkits/computation-FP16/nvidia/gemm.cu | 4 ++++ base/toolkits/computation-FP32/nvidia/gemm.cu | 4 ++++ base/toolkits/computation-FP64/nvidia/gemm.cu | 3 +++ base/toolkits/computation-INT8/nvidia/gemm.cu | 4 ++++ base/toolkits/computation-TF32/nvidia/gemm.cu | 4 ++++ base/toolkits/main_memory-bandwidth/nvidia/bandwidth.cu | 4 ++++ inference/requirements.txt | 1 + 16 files changed, 69 insertions(+), 1 deletion(-) create mode 100644 inference/requirements.txt diff --git a/base/benchmarks/computation-BF16/main.py b/base/benchmarks/computation-BF16/main.py index 9a0afeebd..1bcf255f8 100644 --- a/base/benchmarks/computation-BF16/main.py +++ b/base/benchmarks/computation-BF16/main.py @@ -1,3 +1,8 @@ +# Copyright (c) 2024 BAAI. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License") +#!/usr/bin/env python3 +# -*- coding: UTF-8 -*- import torch import torch.distributed as dist import os diff --git a/base/benchmarks/computation-FP16/main.py b/base/benchmarks/computation-FP16/main.py index 452851b06..5ceefca7e 100644 --- a/base/benchmarks/computation-FP16/main.py +++ b/base/benchmarks/computation-FP16/main.py @@ -1,3 +1,8 @@ +# Copyright (c) 2024 BAAI. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License") +#!/usr/bin/env python3 +# -*- coding: UTF-8 -*- import torch import torch.distributed as dist import os diff --git a/base/benchmarks/computation-FP32/main.py b/base/benchmarks/computation-FP32/main.py index 36ed84e53..cd2161fd5 100644 --- a/base/benchmarks/computation-FP32/main.py +++ b/base/benchmarks/computation-FP32/main.py @@ -1,3 +1,8 @@ +# Copyright (c) 2024 BAAI. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License") +#!/usr/bin/env python3 +# -*- coding: UTF-8 -*- import torch import torch.distributed as dist import os diff --git a/base/benchmarks/computation-FP64/main.py b/base/benchmarks/computation-FP64/main.py index edd41cf88..d25f4a7c9 100644 --- a/base/benchmarks/computation-FP64/main.py +++ b/base/benchmarks/computation-FP64/main.py @@ -1,3 +1,8 @@ +# Copyright (c) 2024 BAAI. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License") +#!/usr/bin/env python3 +# -*- coding: UTF-8 -*- import torch import torch.distributed as dist import os diff --git a/base/benchmarks/computation-TF32/main.py b/base/benchmarks/computation-TF32/main.py index f950d12dd..c94cd588c 100644 --- a/base/benchmarks/computation-TF32/main.py +++ b/base/benchmarks/computation-TF32/main.py @@ -1,3 +1,8 @@ +# Copyright (c) 2024 BAAI. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License") +#!/usr/bin/env python3 +# -*- coding: UTF-8 -*- import torch import torch.distributed as dist import os diff --git a/base/benchmarks/drivers/utils.py b/base/benchmarks/drivers/utils.py index 256b346a3..dfb9b31c9 100644 --- a/base/benchmarks/drivers/utils.py +++ b/base/benchmarks/drivers/utils.py @@ -1,3 +1,8 @@ +# Copyright (c) 2024 BAAI. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License") +#!/usr/bin/env python3 +# -*- coding: UTF-8 -*- import torch @@ -36,4 +41,4 @@ def get_memory_capacity(vendor, rank): return torch.cuda.get_device_properties(rank).total_memory else: print("unspecified vendor {}, return -1.0".format(vendor)) - return -1.0 \ No newline at end of file + return -1.0 diff --git a/base/benchmarks/main_memory-bandwidth/main.py b/base/benchmarks/main_memory-bandwidth/main.py index 3cea75f9c..8b6830cd7 100644 --- a/base/benchmarks/main_memory-bandwidth/main.py +++ b/base/benchmarks/main_memory-bandwidth/main.py @@ -1,3 +1,8 @@ +# Copyright (c) 2024 BAAI. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License") +#!/usr/bin/env python3 +# -*- coding: UTF-8 -*- import torch import torch.distributed as dist import os diff --git a/base/container_main.py b/base/container_main.py index 51bd2665f..629988653 100644 --- a/base/container_main.py +++ b/base/container_main.py @@ -1,3 +1,8 @@ +# Copyright (c) 2024 BAAI. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License") +#!/usr/bin/env python3 +# -*- coding: UTF-8 -*- import time from loguru import logger import os diff --git a/base/toolkits/computation-BF16/nvidia/gemm.cu b/base/toolkits/computation-BF16/nvidia/gemm.cu index 400b6bca6..43fb61441 100644 --- a/base/toolkits/computation-BF16/nvidia/gemm.cu +++ b/base/toolkits/computation-BF16/nvidia/gemm.cu @@ -1,3 +1,7 @@ +// Copyright (c) 2024 BAAI. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License") + #include #include #include diff --git a/base/toolkits/computation-FP16/nvidia/gemm.cu b/base/toolkits/computation-FP16/nvidia/gemm.cu index e1d7f7e0a..69b5da611 100644 --- a/base/toolkits/computation-FP16/nvidia/gemm.cu +++ b/base/toolkits/computation-FP16/nvidia/gemm.cu @@ -1,3 +1,7 @@ +// Copyright (c) 2024 BAAI. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License") + #include #include #include diff --git a/base/toolkits/computation-FP32/nvidia/gemm.cu b/base/toolkits/computation-FP32/nvidia/gemm.cu index 2c44e60fb..6c6e1a27a 100644 --- a/base/toolkits/computation-FP32/nvidia/gemm.cu +++ b/base/toolkits/computation-FP32/nvidia/gemm.cu @@ -1,3 +1,7 @@ +// Copyright (c) 2024 BAAI. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License") + #include #include #include diff --git a/base/toolkits/computation-FP64/nvidia/gemm.cu b/base/toolkits/computation-FP64/nvidia/gemm.cu index d3228c6d5..78e0b49a6 100644 --- a/base/toolkits/computation-FP64/nvidia/gemm.cu +++ b/base/toolkits/computation-FP64/nvidia/gemm.cu @@ -1,3 +1,6 @@ +// Copyright (c) 2024 BAAI. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License") #include #include #include diff --git a/base/toolkits/computation-INT8/nvidia/gemm.cu b/base/toolkits/computation-INT8/nvidia/gemm.cu index 3571b1408..ec179a9a9 100644 --- a/base/toolkits/computation-INT8/nvidia/gemm.cu +++ b/base/toolkits/computation-INT8/nvidia/gemm.cu @@ -1,3 +1,7 @@ +// Copyright (c) 2024 BAAI. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License") + #include #include #include diff --git a/base/toolkits/computation-TF32/nvidia/gemm.cu b/base/toolkits/computation-TF32/nvidia/gemm.cu index 60c79b7f4..eeadf82bc 100644 --- a/base/toolkits/computation-TF32/nvidia/gemm.cu +++ b/base/toolkits/computation-TF32/nvidia/gemm.cu @@ -1,3 +1,7 @@ +// Copyright (c) 2024 BAAI. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License") + #include #include #include diff --git a/base/toolkits/main_memory-bandwidth/nvidia/bandwidth.cu b/base/toolkits/main_memory-bandwidth/nvidia/bandwidth.cu index dd5e1a0ad..deb81dd3c 100644 --- a/base/toolkits/main_memory-bandwidth/nvidia/bandwidth.cu +++ b/base/toolkits/main_memory-bandwidth/nvidia/bandwidth.cu @@ -1,3 +1,7 @@ +// Copyright (c) 2024 BAAI. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License") + #include #include diff --git a/inference/requirements.txt b/inference/requirements.txt new file mode 100644 index 000000000..8e6ceeab2 --- /dev/null +++ b/inference/requirements.txt @@ -0,0 +1 @@ +loguru==0.7.2 From 0fc654754bb9cc888bff4555da8baee243930b5e Mon Sep 17 00:00:00 2001 From: shh2000 <13820618441@163.com> Date: Sun, 28 Apr 2024 11:45:35 +0800 Subject: [PATCH 4/7] replace --- training/run_benchmarks/config/test_conf.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/training/run_benchmarks/config/test_conf.py b/training/run_benchmarks/config/test_conf.py index 1332f29c8..84887a447 100644 --- a/training/run_benchmarks/config/test_conf.py +++ b/training/run_benchmarks/config/test_conf.py @@ -62,12 +62,13 @@ ''' CASES = { # nvidia cases - "bert:pytorch_1.8:A100:1:8:1": "/raid/home_datasets_ckpt/bert/train/", - "glm:pytorch_1.8:A100:1:8:1": "/raid/home_datasets_ckpt/glm/train/", - "cpm:pytorch_1.8:A100:1:8:1": "/raid/home_datasets_ckpt/cpm/train/", + "llama3_8b:megatron_core060:A100:1:8:1": "/data/llama3_8b_pretrain" + # "bert:pytorch_1.8:A100:1:8:1": "/raid/home_datasets_ckpt/bert/train/", + # "glm:pytorch_1.8:A100:1:8:1": "/raid/home_datasets_ckpt/glm/train/", + # "cpm:pytorch_1.8:A100:1:8:1": "/raid/home_datasets_ckpt/cpm/train/", - #"llama2_7b_finetune:pytorch_2.0.1:A100:1:1:1": "/raid/dataset/llama2_finetune/", - #"aquila2_7b_finetune:flagscale:A800:1:8:1": "/raid/dataset/aquila2_7b_finetune", + # "llama2_7b_finetune:pytorch_2.0.1:A100:1:1:1": "/raid/dataset/llama2_finetune/", + # "aquila2_7b_finetune:flagscale:A800:1:8:1": "/raid/dataset/aquila2_7b_finetune", # "mobilenetv2:pytorch_1.8:A100:1:8:1": "/raid/dataset/ImageNet_1k_2012/", # "vit:pytorch_1.13:A100:1:8:1": "/raid/dataset/ImageNet_1k_2012/", # "efficientnet:pytorch_1.13:A100:1:8:1": "/raid/dataset/ImageNet_1k_2012/", From 0f2c13b7cad5c0ace534c033f4651e45ef286beb Mon Sep 17 00:00:00 2001 From: shh2000 <13820618441@163.com> Date: Sun, 28 Apr 2024 16:06:31 +0800 Subject: [PATCH 5/7] ur --- training/benchmarks/llama3_8B/megatron/README.md | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/training/benchmarks/llama3_8B/megatron/README.md b/training/benchmarks/llama3_8B/megatron/README.md index ed07a6ef4..a63f3d0dc 100644 --- a/training/benchmarks/llama3_8B/megatron/README.md +++ b/training/benchmarks/llama3_8B/megatron/README.md @@ -8,23 +8,29 @@ LLaMA3 is a new generation of large language models developed by Meta. The first - 模型代码来源 -This case includes code from the LLAMA 3 COMMUNITY LICENSE AGREEMENT License open source project at:[meta-llama/Meta-Llama-3-8B · Hugging Face](https://huggingface.co/meta-llama/Meta-Llama-3-8B) +我们假设您了解Meta开源了Llama3的哪些部分（例如包括模型权重文件、tokenizer等，不包括预训练代码、预训练实现、预训练数据集等）。出于上述事实，本评测样例基于开源Megatron框架，使用开源wudao数据集，在[meta-llama/Meta-Llama-3-8B · Hugging Face](https://huggingface.co/meta-llama/Meta-Llama-3-8B)设计的LLaMA3-8B算法结构上进行预训练，来进行AI硬件评测。测试样例内部代码为FlagPerf编写。需要下载或准备的文件见**数据准备**小节，依赖的外部软件或信息见**依赖**小节。 # 数据准备 ### 模型配置及tokenizer准备 -本测试样例为预训练case，需要下载tokenizer，不需要下载模型代码或模型权重。tokenizer需向llama3官方申请并下载8B版本，并在data_dir下创建\目录，按照[unsloth/llama-3-8b at main (huggingface.co)](https://huggingface.co/unsloth/llama-3-8b/tree/main)仓库示例格式存放（不需要权重部分）。其中，\目录可更改，需在形如config\_A100\_1x8.py的配置文件中同步修改。默认为\。 +本测试样例为预训练case，需要下载tokenizer（我们假设您了解tokenizer包括哪些文件。如不确定，可下载llama3-8B所有文件，尽管我们不需要使用模型权重文件），不需要下载模型代码或模型权重。tokenizer需向llama3官方申请并下载8B所用版本，并在data_dir下创建llama3_8b_hf目录，按照huggingface要求的格式进行处理或存放。了解data\_dir需要阅读FlagPerf有关训练的文档，或不阅读相关文档直接修改FlagPerf/training/run_benchmarks/config/test_conf.py中CASES变量中的value。 + +出于对Llama3开源协议的遵守，尽管不推荐，不了解相关背景的用户仍可以参考[unsloth/llama-3-8b at main (huggingface.co)](https://huggingface.co/unsloth/llama-3-8b/tree/main)仓库示例格式存放tokenizer相关文件（不需要模型权重部分，如果您不了解哪些是模型权重部分，请保留全部文件）。其中，data_dir下面创建的llama3_8b_hf目录名称可更改。如更改，需在形如FlagPerf/training/nvidia/llama3\_8B-megatron/config/config\_A100\_1x8.py的配置文件中同步修改。默认为llama3\_8b_hf。 ### 数据集准备 本测试样例数据使用智源研究院wudao数据集，需向llama3官方申请tokenizer并进行预处理。测试数据集的内容不是影响AI硬件评测结果的核心因素，可以参考或使用阿里灵骏团队开放的预处理好的数据集[Pai-Megatron-Patch/examples/llama3/README.md at main · alibaba/Pai-Megatron-Patch (github.com)](https://github.com/alibaba/Pai-Megatron-Patch/blob/main/examples/llama3/README.md#%E6%95%B0%E6%8D%AE%E9%9B%86%E5%92%8C%E6%A8%A1%E5%9E%8B%E4%B8%8B%E8%BD%BD) -将上述两个文件（.bin与.idx，不清楚文件用法可以参考Megatron-LM仓库相关介绍）放置于data_dir下。 +在上述README文件中找到并执行 + +`wget https://atp-modelzoo-wlcb-pai.oss-cn-wulanchabu.aliyuncs.com/release/models/pai-megatron-patch/llama3-datasets/wudao_llama3bpe_content_document.binwget https://atp-modelzoo-wlcb-pai.oss-cn-wulanchabu.aliyuncs.com/release/models/pai-megatron-patch/llama3-datasets/wudao_llama3bpe_content_document.idx` + +将上述两个文件（.bin与.idx，不清楚文件原理可以参考Megatron-LM仓库相关介绍）放置于data_dir下。 -# 相关开源仓库 +# 依赖 ### llama3（算法结构、tokenizer） From ad40d234a2862f94efc3f2173e1ebcc5d7a5eb8f Mon Sep 17 00:00:00 2001 From: shh2000 <13820618441@163.com> Date: Sun, 28 Apr 2024 16:08:41 +0800 Subject: [PATCH 6/7] ur --- training/benchmarks/llama3_8B/megatron/README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/training/benchmarks/llama3_8B/megatron/README.md b/training/benchmarks/llama3_8B/megatron/README.md index a63f3d0dc..2ce29286c 100644 --- a/training/benchmarks/llama3_8B/megatron/README.md +++ b/training/benchmarks/llama3_8B/megatron/README.md @@ -8,7 +8,7 @@ LLaMA3 is a new generation of large language models developed by Meta. The first - 模型代码来源 -我们假设您了解Meta开源了Llama3的哪些部分（例如包括模型权重文件、tokenizer等，不包括预训练代码、预训练实现、预训练数据集等）。出于上述事实，本评测样例基于开源Megatron框架，使用开源wudao数据集，在[meta-llama/Meta-Llama-3-8B · Hugging Face](https://huggingface.co/meta-llama/Meta-Llama-3-8B)设计的LLaMA3-8B算法结构上进行预训练，来进行AI硬件评测。测试样例内部代码为FlagPerf编写。需要下载或准备的文件见**数据准备**小节，依赖的外部软件或信息见**依赖**小节。 +我们假设您了解Meta开源了Llama3的哪些部分（例如包括模型权重文件、tokenizer等，不包括预训练代码、预训练实现、预训练数据集等）。出于上述事实，本评测样例基于开源Megatron框架，使用开源wudao数据集，在[meta-llama/Meta-Llama-3-8B · Hugging Face](https://huggingface.co/meta-llama/Meta-Llama-3-8B)设计的LLaMA3-8B算法结构上进行预训练，来进行AI硬件评测。测试样例代码为FlagPerf编写。需要下载或准备的文件见**数据准备**小节，依赖的外部软件或信息见**依赖**小节。 # 数据准备 @@ -49,4 +49,4 @@ iii. You must retain in all copies of the Llama Materials that you distribute th ### megatron-core（计算任务实现） -[NVIDIA/Megatron-LM: Ongoing research training transformer models at scale (github.com)](https://github.com/NVIDIA/Megatron-LM/tree/main?tab=License-1-ov-file#readme) \ No newline at end of file +[NVIDIA/Megatron-LM: Ongoing research training transformer models at scale (github.com)](https://github.com/NVIDIA/Megatron-LM/tree/main?tab=License-1-ov-file#readme) From a188d3e803fbfe4d13e6a7401e1ec917c99e44fc Mon Sep 17 00:00:00 2001 From: shh2000 <13820618441@163.com> Date: Sun, 28 Apr 2024 17:39:36 +0800 Subject: [PATCH 7/7] ur --- training/benchmarks/llama3_8B/megatron/README.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/training/benchmarks/llama3_8B/megatron/README.md b/training/benchmarks/llama3_8B/megatron/README.md index 2ce29286c..b9f6da3c4 100644 --- a/training/benchmarks/llama3_8B/megatron/README.md +++ b/training/benchmarks/llama3_8B/megatron/README.md @@ -8,16 +8,16 @@ LLaMA3 is a new generation of large language models developed by Meta. The first - 模型代码来源 -我们假设您了解Meta开源了Llama3的哪些部分（例如包括模型权重文件、tokenizer等，不包括预训练代码、预训练实现、预训练数据集等）。出于上述事实，本评测样例基于开源Megatron框架，使用开源wudao数据集，在[meta-llama/Meta-Llama-3-8B · Hugging Face](https://huggingface.co/meta-llama/Meta-Llama-3-8B)设计的LLaMA3-8B算法结构上进行预训练，来进行AI硬件评测。测试样例代码为FlagPerf编写。需要下载或准备的文件见**数据准备**小节，依赖的外部软件或信息见**依赖**小节。 +Meta仅开源了Llama3的模型权重文件、tokenizer等，不包括预训练代码、预训练实现、预训练数据集等。出于上述事实，本评测样例基于开源Megatron框架，使用开源wudao数据集，在[meta-llama/Meta-Llama-3-8B · Hugging Face](https://huggingface.co/meta-llama/Meta-Llama-3-8B)设计的LLaMA3-8B算法结构上进行预训练，来进行AI硬件评测。测试样例代码为FlagPerf编写。需要下载或准备的文件见**数据准备**小节，依赖的外部软件或信息见**依赖**小节。 # 数据准备 ### 模型配置及tokenizer准备 -本测试样例为预训练case，需要下载tokenizer（我们假设您了解tokenizer包括哪些文件。如不确定，可下载llama3-8B所有文件，尽管我们不需要使用模型权重文件），不需要下载模型代码或模型权重。tokenizer需向llama3官方申请并下载8B所用版本，并在data_dir下创建llama3_8b_hf目录，按照huggingface要求的格式进行处理或存放。了解data\_dir需要阅读FlagPerf有关训练的文档，或不阅读相关文档直接修改FlagPerf/training/run_benchmarks/config/test_conf.py中CASES变量中的value。 +本测试样例为预训练case，需要下载tokenizer（如不确定tokenizer包含哪些文件，可下载llama3-8B所有文件，尽管我们不需要使用模型权重文件），不需要下载模型代码或模型权重。tokenizer需向llama3官方申请并下载8B所用版本，并在data_dir下创建llama3_8b_hf目录，按照huggingface要求的格式进行处理或存放。了解data\_dir需要阅读FlagPerf有关训练的文档，或直接修改FlagPerf/training/run_benchmarks/config/test_conf.py中CASES变量中的value。 -出于对Llama3开源协议的遵守，尽管不推荐，不了解相关背景的用户仍可以参考[unsloth/llama-3-8b at main (huggingface.co)](https://huggingface.co/unsloth/llama-3-8b/tree/main)仓库示例格式存放tokenizer相关文件（不需要模型权重部分，如果您不了解哪些是模型权重部分，请保留全部文件）。其中，data_dir下面创建的llama3_8b_hf目录名称可更改。如更改，需在形如FlagPerf/training/nvidia/llama3\_8B-megatron/config/config\_A100\_1x8.py的配置文件中同步修改。默认为llama3\_8b_hf。 +出于对Llama3开源协议的遵守，尽管不推荐，不了解相关背景的用户仍可以参考[unsloth/llama-3-8b at main (huggingface.co)](https://huggingface.co/unsloth/llama-3-8b/tree/main)仓库示例格式存放tokenizer相关文件（不需要模型权重部分，如不确定，请保留全部文件）。其中，data_dir下面创建的llama3_8b_hf目录名称可更改。如更改，需在形如FlagPerf/training/nvidia/llama3\_8B-megatron/config/config\_A100\_1x8.py的配置文件中同步修改。默认为llama3\_8b_hf。 ### 数据集准备