Moxin-LLM

article

#LICENSE

Apache License

Version 2.0 , January 2004

http://www. apache. org/licenses/

TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION

Definitions. "License" shall mean the terms and conditions for use, reproduction,

and distribution as defined by Sections 1 through 9 of this document. "Licensor" shall mean the copyright owner or entity authorized by

the copyright owner that is granting the License. "Legal Entity" shall mean the union of the acting entity and all

other entities that control, are controlled by, or are under common

control with that entity. For the purposes of this definition,

"control" means (i) the power, direct or indirect, to cause the

direction or management of such entity, whether by contract or

otherwise, or (ii) ownership of fifty percent (50%) or more of the

outstanding shares, or (iii) beneficial ownership of such entity. "You" (or "Your") shall mean an individual or Legal Entity

exercising permissions granted by this License. "Source" form shall mean the preferred form for making modifications,

including but not limited to software source code, documentation

source, and configuration files. "Object" form shall mean any form resulting from mechanical

transformation or translation of a Source form, including but

not limited to compiled object code, generated documentation,

and conversions to other media types. "Work" shall mean the work of authorship, whether in Source or

Object form, made available under the License, as indicated by a

(an example is provided in the Appendix below).

"Derivative Works" shall mean any work, whether in Source or Object

form, that is based on (or derived from) the Work and for which the

editorial revisions, annotations, elaborations, or other modifications

represent, as a whole, an original work of authorship. For the purposes

of this License, Derivative Works shall not include works that remain

separable from, or merely link (or bind by name) to the interfaces of,

the Work and Derivative Works thereof. "Contribution" shall mean any work of authorship, including

the original version of the Work and any modifications or additions

to that Work or Derivative Works thereof, that is intentionally

submitted to Licensor for inclusion in the Work by the copyright owner

or by an individual or Legal Entity authorized to submit on behalf of

the copyright owner. For the purposes of this definition, "submitted"

means any form of electronic, verbal, or written communication sent

to the Licensor or its representatives, including but not limited to

communication on electronic mailing lists, source code control systems,

and issue tracking systems that are managed by, or on behalf of, the

Licensor for the purpose of discussing and improving the Work, but

excluding communication that is conspicuously marked or otherwise

designated in writing by the copyright owner as "Not a Contribution."

"Contributor" shall mean Licensor and any individual or Legal Entity

on behalf of whom a Contribution has been received by Licensor and

subsequently incorporated within the Work. 2.

Grant of Copyright License. Subject to the terms and conditions of

this License, each Contributor hereby grants to You a perpetual,

worldwide, non-exclusive, no-charge, royalty-free, irrevocable

publicly display, publicly perform, sublicense, and distribute the

Work and such Derivative Works in Source or Object form. 3.

Grant of Patent License. Subject to the terms and conditions of

this License, each Contributor hereby grants to You a perpetual,

worldwide, non-exclusive, no-charge, royalty-free, irrevocable

(except as stated in this section) patent license to make, have made,

use, offer to sell, sell, import, and otherwise transfer the Work,

where such license applies only to those patent claims licensable

by such Contributor that are necessarily infringed by their

Contribution(s) alone or by combination of their Contribution(s)

with the Work to which such Contribution(s) was submitted. If You

institute patent litigation against any entity (including a

cross-claim or counterclaim in a lawsuit) alleging that the Work

or a Contribution incorporated within the Work constitutes direct

or contributory patent infringement, then any patent licenses

granted to You under this License for that Work shall terminate

as of the date such litigation is filed. 4.

Redistribution. You may reproduce and distribute copies of the

Work or Derivative Works thereof in any medium, with or without

modifications, and in Source or Object form, provided that You

meet the following conditions:

(a) You must give any other recipients of the Work or

Derivative Works a copy of this License; and

(b) You must cause any modified files to carry prominent notices

stating that You changed the files; and

that You distribute, all copyright, patent, trademark, and

attribution notices from the Source form of the Work,

excluding those notices that do not pertain to any part of

the Derivative Works; and

(d) If the Work includes a "NOTICE" text file as part of its

distribution, then any Derivative Works that You distribute must

include a readable copy of the attribution notices contained

within such NOTICE file, excluding those notices that do not

pertain to any part of the Derivative Works, in at least one

of the following places: within a NOTICE text file distributed

as part of the Derivative Works; within the Source form or

documentation, if provided along with the Derivative Works; or,

within a display generated by the Derivative Works, if and

wherever such third-party notices normally appear. The contents

of the NOTICE file are for informational purposes only and

do not modify the License. You may add Your own attribution

notices within Derivative Works that You distribute, alongside

or as an addendum to the NOTICE text from the Work, provided

that such additional attribution notices cannot be construed

as modifying the License. You may add Your own copyright statement to Your modifications and

may provide additional or different license terms and conditions

for use, reproduction, or distribution of Your modifications, or

for any such Derivative Works as a whole, provided Your use,

reproduction, and distribution of the Work otherwise complies with

the conditions stated in this License. 5.

Submission of Contributions. Unless You explicitly state otherwise,

any Contribution intentionally submitted for inclusion in the Work

by You to the Licensor shall be under the terms and conditions of

this License, without any additional terms or conditions. Notwithstanding the above, nothing herein shall supersede or modify

the terms of any separate license agreement you may have executed

with Licensor regarding such Contributions. 6.

Trademarks. This License does not grant permission to use the trade

names, trademarks, service marks, or product names of the Licensor,

except as required for reasonable and customary use in describing the

origin of the Work and reproducing the content of the NOTICE file. 7.

Disclaimer of Warranty. Unless required by applicable law or

agreed to in writing, Licensor provides the Work (and each

Contributor provides its Contributions) on an "AS IS" BASIS,

WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or

implied, including, without limitation, any warranties or conditions

of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A

PARTICULAR PURPOSE. You are solely responsible for determining the

appropriateness of using or redistributing the Work and assume any

risks associated with Your exercise of permissions under this License. 8.

Limitation of Liability. In no event and under no legal theory,

whether in tort (including negligence), contract, or otherwise,

unless required by applicable law (such as deliberate and grossly

negligent acts) or agreed to in writing, shall any Contributor be

liable to You for damages, including any direct, indirect, special,

incidental, or consequential damages of any character arising as a

result of this License or out of the use or inability to use the

Work (including but not limited to damages for loss of goodwill,

work stoppage, computer failure or malfunction, or any and all

other commercial damages or losses), even if such Contributor

has been advised of the possibility of such damages. 9.

Accepting Warranty or Additional Liability. While redistributing

the Work or Derivative Works thereof, You may choose to offer,

and charge a fee for, acceptance of support, warranty, indemnity,

or other liability obligations and/or rights consistent with this

License. However, in accepting such obligations, You may act only

on Your own behalf and on Your sole responsibility, not on behalf

of any other Contributor, and only if You agree to indemnify,

defend, and hold each Contributor harmless for any liability

incurred by, or claims asserted against, such Contributor by reason

of your accepting any such warranty or additional liability. END OF TERMS AND CONDITIONS

APPENDIX: How to apply the Apache License to your work. To apply the Apache License to your work, attach the following

boilerplate notice, with the fields enclosed by brackets "[]"

replaced with your own identifying information. (Don't include

the brackets!) The text should be enclosed in the appropriate

comment syntax for the file format. We also recommend that a

file or class name and description of purpose be included on the

same "printed page" as the copyright notice for easier

Licensed under the Apache License, Version 2.0 (the "License");

you may not use this file except in compliance with the License. You may obtain a copy of the License at

http://www. apache. org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software

distributed under the License is distributed on an "AS IS" BASIS,

WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and

limitations under the License. \n[SEP]\n#REQUIREMENTS. md \n

<h1 align="center"> Moxin LLM</h1>

<p align="center"> <a href="https://arxiv. org /abs/2412.06845 ">Technical Report</a> &nbsp&nbsp | &nbsp&nbsp <a href="https://huggingface. co /moxin-org/moxin-llm-7b">Base Model</a> &nbsp&nbsp | &nbsp&nbsp <a href="https://huggingface. co /moxin-org/moxin-chat-7b">Chat Model</a> </p>

## Introduction

Generative AI (GAI) offers unprecedented opportunities for research and innovation, but its commercialization has raised concerns about transparency, reproducibility, and safety. Many open GAI models lack the necessary components for full understanding and reproducibility, and some use restrictive licenses whilst claiming to be “open-source”.

To address these concerns, we follow the [Model Openness Framework (MOF)](https://arxiv. org /pdf/2403.13784 ), a ranked classification system that rates machine learning models based on their completeness and openness, following principles of open science, open source, open data, and open access. By promoting transparency and reproducibility, the MOF combats “openwashing” practices and establishes completeness and openness as primary criteria alongside the core tenets of responsible AI. Wide adoption of the MOF will foster a more open AI ecosystem, benefiting research, innovation, and adoption of state-of-the-art models. We follow MOF to release the datasets during training, the training scripts, and the trained models. ## Model

You can download our base 7B model from this [link](https://huggingface. co /moxin-org/moxin-llm-7b) and our chat 7B model from this [link](https://huggingface. co /moxin-org/moxin-chat-7b).

## Evaluation

We test the performance of our model with [lm-evaluation-harness](https://github. com /EleutherAI/lm-evaluation-harness).

The evaluation results on common datasets are shown below. We test on AI2 Reasoning Challenge (25-shot), HellaSwag (10-shot), MMLU (5-shot), and Winogrande (5-shot).

|:----------------------:|:-----:|:---------:|:-----:|:---------:|:-----:|

| Mistral-7B | 57.59 | 83.25 | 62.42 | 78.77 | 70.51 |

| LLaMA 3.1 -8B | 54.61 | 81.95 | 65.16 | 77.35 | 69.77 |

| LLaMA 3-8B | 55.46 | 82.09 | 65.29 | 77.82 | 70.17 |

| LLaMA 2-7B | 49.74 | 78.94 | 45.89 | 74.27 | 62.21 |

| Qwen 2-7B | 57.68 | 80.76 | 70.42 | 77.43 | 71.57 |

| gemma-7b | 56.48 | 82.31 | 63.02 | 78.3 | 70.03 |

| internlm2.5 -7b | 54.78 | 79.7 | 68.17 | 80.9 | 70.89 |

| Baichuan2-7B | 47.87 | 73.89 | 54.13 | 70.8 | 61.67 |

| Yi-1.5 -9B | 58.36 | 80.36 | 69.54 | 77.53 | 71.48 |

| Moxin-7B-original | 53.75 | 75.46 | 59.43 | 70.32 | 64.74 |

| Moxin-7B-finetuned | 59.47 | 83.08 | 60.97 | 78.69 | 70.55 |

We also test the zero shot performance on AI2 Reasoning Challenge (0-shot), AI2 Reasoning Easy (0-shot), HellaSwag (0-shot), PIQA (0-shot) and Winogrande (0-shot).

|:-----------------: |:---------: |:---------: |:-----: |:-----: |:-----: |:-----: |

| Mistral-7B | 80.39 | 73.4 | 82.15 | 78.28 | 52.22 | 73.29 |

| LLaMA 2-7B | 75.99 | 69.06 | 79.11 | 74.54 | 46.42 | 69.02 |

| LLaMA 2-13B | 79.37 | 72.22 | 80.52 | 77.4 | 49.06 | 71.71 |

| LLaMA 3.1 -8B | 78.92 | 74.19 | 81.12 | 81.06 | 53.67 | 73.79 |

| gemma-7b | 80.45 | 73.72 | 80.9 | 79.97 | 54.1 | 73.83 |

| Qwen v2-7B | 78.9 | 72.38 | 79.98 | 74.71 | 50.09 | 71.21 |

| internlm2.5 -7b | 79.14 | 77.9 | 80.52 | 76.16 | 51.37 | 73.02 |

| Baichuan2-7B | 72.25 | 67.17 | 77.26 | 72.98 | 42.15 | 66.36 |

| Yi-1.5 -9B | 77.86 | 73.01 | 80.74 | 79.04 | 55.03 | 73.14 |

| deepseek-7b | 76.13 | 69.77 | 79.76 | 71.04 | 44.8 | 68.3 |

| Moxin-7B-original | 72.06 | 66.31 | 78.07 | 71.47 | 48.15 | 67.21 |

| Moxin-7B-finetune | 80.03 | 75.17 | 82.24 | 81.12 | 58.64 | 75.44 |

## Inference

You can use the following code to run inference with the model. The model is saved under './model/' directory. Change the model directory accordingly or use the Huggingface link. ```

import torch

from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

torch. backends. cuda. enable_mem_efficient_sdp (False)

torch. backends. cuda. enable_flash_sdp (False)

model_name = 'moxin-org/moxin-7b'

tokenizer = AutoTokenizer. from_pretrained (model_name)

model = AutoModelForCausalLM. from_pretrained (

model_name,

torch_dtype=torch. bfloat16 ,

device_map="auto",

trust_remote_code=True,

)

pipe = pipeline(

"text-generation",

model=model,

tokenizer = tokenizer,

torch_dtype=torch. bfloat16 ,

device_map="auto"

)

prompt = "Can you explain the concept of regularization in machine learning?"

sequences = pipe(

prompt,

do_sample=True,

max_new_tokens=100,

temperature=0.7 ,

top_k=50,

top_p=0.95 ,

num_return_sequences=1,

)

print(sequences[0]['generated_text'])

```

## Environment

### 1.

Dataset config

To prepare the dataset, it needs to install the following package,

```

pip install datasets

```

### 2.

Cuda install

We use cuda 11.7.

Other cuda versions may also work. ```

get https://developer. download. nvidia. com /compute/cuda/11.7.0/local_installers/cuda_11.7.0_515.43.04_linux. run

sudo sh cuda_11.7.0_515.43.04_linux. run

```

### 3.

Install pytorch

We use pytorch 2.0.0.

```

conda create --name llm_train python==3.10

conda activate llm_train

pip install torch==2.0.0 torchvision==0.15.1 torchaudio==2.0.1

```

### 4.

Install other packages

To install other packages, follow the requirements. txt

```

pip install -r requirements. txt

```

### 5.

Install flash attention

We use flash-attention 2.2.1.

```

git clone https://github. com /Dao-AILab/flash-attention. git

cd flash-attention/

git checkout a1576ad ## flash-attention 2.2.1

python setup. py install

cd./csrc

cd fused_dense_lib && pip install -v. cd../xentropy && pip install -v. cd../rotary && pip install -v. cd../layer_norm && pip install -v. ```

## Datasets

To use the [SlimPajama dataset](https://huggingface. co /datasets/cerebras/SlimPajama-627B) for pretraining, you can download the dataset using Hugging Face datasets:

```

import datasets

ds = datasets. load_dataset ("cerebras/SlimPajama-627B")

```

SlimPajama is the largest extensively deduplicated, multi-corpora, open-source dataset for training large language models. SlimPajama was created by cleaning and deduplicating the 1.2T token RedPajama dataset from Together. By filtering out low quality data and duplicates, it removes 49.6 % of bytes, slimming down the RedPajama dataset from 1210B to 627B tokens. SlimPajama offers the highest quality and most compute efficient data to train on for runs up to 627B tokens. When upsampled, SlimPajama is expected to perform equal to or better than RedPajama-1T when training at trillion token scale. To use the [stack-dedup dataset](https://huggingface. co /datasets/bigcode/the-stack-dedup) for pretraining, you can download the dataset using Hugging Face datasets:

```

from datasets import load_dataset

# full dataset (3TB of data)

ds = load_dataset("bigcode/the-stack-dedup", split="train")

# specific language (e. g. Dockerfiles)

ds = load_dataset("bigcode/the-stack-dedup", data_dir="data/dockerfile", split="train")

# dataset streaming (will only download the data as needed)

ds = load_dataset("bigcode/the-stack-dedup", streaming=True, split="train")

for sample in iter(ds): print(sample["content"])

```

The Stack contains over 6TB of permissively-licensed source code files covering 358 programming languages. The dataset was created as part of the BigCode Project, an open scientific collaboration working on the responsible development of Large Language Models for Code (Code LLMs).

The Stack serves as a pre-training dataset for Code LLMs, i. e., code-generating AI systems which enable the synthesis of programs from natural language descriptions as well as other from code snippets. This is the near-deduplicated version with 3TB data. You can find more details about the DCLM-baseline dataset on the [homepage](https://huggingface. co /datasets/mlfoundations/dclm-baseline-1.0 ).

## Training

We follow the [ColossalAI](https://github. com /hpcaitech/ColossalAI) framework to train the LLM model. Colossal-AI provides a collection of parallel components for the training. It aims to support to write the distributed deep learning models just like how you write your model on your laptop. It provides user-friendly tools to kickstart distributed training and inference in a few lines. We provide a few examples to show how to run benchmark or pretraining based on Colossal-AI. ### 1.

Training LLM

You can find the shell scripts in 'scripts/train_7B' directory. The main command should be in the format of:

```

colossalai run --nproc_per_node YOUR_GPU_PER_NODE --hostfile YOUR_HOST_FILE \

benchmark. py --OTHER_CONFIGURATIONS

```

#### a. Running on a sinlge node

we provide an example to run the training on a single node as below,

```

colossalai run --nproc_per_node 1 pretrain. py \

--config 7b \

--dataset togethercomputer/RedPajama-Data-1T-Sample \

--batch_size 1 \

--num_epochs 5 \

--save_interval 5000 \

--max_length 2048 \

--save_dir output-checkpoints \

--plugin zero2_cpu \

--lr 2e-5 \

--expanded_model hpcai-tech/Colossal-LLaMA-2-7b-base

```

In the example, it uses the sample dataset 'togethercomputer/RedPajama-Data-1T-Sample' for training. It trains the 7B model 'hpcai-tech/Colossal-LLaMA-2-7b-base'.

You can refer the main file 'run. sh ' and 'pretrain. py ' for more details. To start the training, run the following,

```bash

bash run. sh

```

#### b. Running on a sinlge node

we provide an example to run the training on multiple nodes as below,

```

srun colossalai run --num_nodes 8 --nproc_per_node 8 pretrain. py \

--config 7b \

--dataset cerebras/SlimPajama-627B \

--batch_size 1 \

--num_epochs 10 \

--save_interval 50000 \

--max_length 2048 \

--save_dir output-checkpoints \

--flash_attention \

--plugin zero2_cpu \

--lr 1e-5 \

--expanded_model hpcai-tech/Colossal-LLaMA-2-7b-base

```

It uses 8 nodes. Put your host file (`hosts. txt `) in this directory with your real host ip or host name. Here is a sample `hosts. txt `:

```text

hostname1

hostname2

hostname3...

hostname8

```

You can refer to the main file 'run-multi-server. sh ' and 'pretrain. py ' for more details. To start the training, run the following,

```bash

bash run-multi-server. sh

```

### 2.

Benchmark

You can find the shell scripts in 'scripts/benchmark_7B' directory. The benchmark mainly test the throughput of the LLM, without actual model training. The main command should be in the format of:

```

colossalai run --nproc_per_node YOUR_GPU_PER_NODE --hostfile YOUR_HOST_FILE \

benchmark. py --OTHER_CONFIGURATIONS

```

Here we will show an example of how to run training llama pretraining with 'gemini, batch_size=16, sequence_length=4096, gradient_checkpoint=True, flash_attn=True'.

#### a. Running environment

This experiment was performed on 4 computing nodes with 32 L40S GPUs in total for LLaMA-2 7B. The nodes are connected with RDMA and GPUs within one node are fully connected with NVLink. #### b. Running command

```bash

cd scripts/benchmark_7B

```

First, put your host file (`hosts. txt `) in this directory with your real host ip or host name. Here is a sample `hosts. txt `:

```text

hostname1

hostname2

hostname3

hostname4

```

Then add environment variables to script if needed. Finally, run the following command to start training:

```bash

bash gemini. sh

```

## Citation

\n[SEP]\n#requirements. txt \n

absl-py==2.1.0

aiohttp==3.9.3

aiosignal==1.3.1

annotated-types==0.6.0

async-timeout==4.0.3

attrs==23.2.0

bcrypt==4.1.2

beautifulsoup4==4.12.3

cachetools==5.3.3

certifi==2024.2.2

cffi==1.16.0

cfgv==3.4.0

charset-normalizer==3.3.2

click==8.1.7

cmake==3.29.0.1

colossalai==0.3.6

contexttimer==0.3.3

cryptography==42.0.5

datasets==2.18.0

decorator==5.1.1

Deprecated==1.2.14

dill==0.3.8

distlib==0.3.8

einops==0.7.0

fabric==3.2.2

filelock==3.13.3

flash-attn==2.2.1

frozenlist==1.4.1

fsspec==2024.2.0

google==3.0.0

google-auth==2.29.0

google-auth-oauthlib==1.0.0

grpcio==1.62.1

huggingface-hub==0.22.2

identify==2.5.35

idna==3.6

invoke==2.2.0

Jinja2==3.1.3

jsonschema==4.21.1

jsonschema-specifications==2023.12.1

lit==18.1.2

Markdown==3.6

markdown-it-py==3.0.0

MarkupSafe==2.1.5

mdurl==0.1.2

mpmath==1.3.0

msgpack==1.0.8

multidict==6.0.5

multiprocess==0.70.16

networkx== 3.3

ninja==1.11.1.1

nodeenv==1.8.0

numpy==1.26.4

nvidia-cublas-cu11==11.10.3.66

nvidia-cublas-cu12==12.1.3.1

nvidia-cuda-cupti-cu11==11.7.101

nvidia-cuda-cupti-cu12==12.1.105

nvidia-cuda-nvrtc-cu11==11.7.99

nvidia-cuda-nvrtc-cu12==12.1.105

nvidia-cuda-runtime-cu11==11.7.99

nvidia-cuda-runtime-cu12==12.1.105

nvidia-cudnn-cu11==8.5.0.96

nvidia-cudnn-cu12==8.9.2.26

nvidia-cufft-cu11==10.9.0.58

nvidia-cufft-cu12==11.0.2.54

nvidia-curand-cu11==10.2.10.91

nvidia-curand-cu12==10.3.2.106

nvidia-cusolver-cu11==11.4.0.1

nvidia-cusolver-cu12==11.4.5.107

nvidia-cusparse-cu11==11.7.4.91

nvidia-cusparse-cu12==12.1.0.106

nvidia-nccl-cu11==2.14.3

nvidia-nccl-cu12==2.19.3

nvidia-nvjitlink-cu12==12.4.127

nvidia-nvtx-cu11==11.7.91

nvidia-nvtx-cu12==12.1.105

oauthlib==3.2.2

packaging==24.0

pandas==2.2.1

paramiko==3.4.0

pip==23.3.1

platformdirs==4.2.0

pre-commit==3.7.0

protobuf==5.26.1

psutil==5.9.8

pyarrow==15.0.2

pyarrow-hotfix==0.6

pyasn1==0.6.0

pyasn1_modules==0.4.0

pycparser==2.22

pydantic==2.6.4

pydantic_core==2.16.3

Pygments==2.17.2

PyNaCl==1.5.0

python-dateutil==2.9.0.post0

pytz==2024.1

PyYAML==6.0.1

ray==2.10.0

referencing==0.34.0

regex==2023.12.25

requests==2.31.0

requests-oauthlib==2.0.0

rich==13.7.1

rpds-py==0.18.0

rsa==4.9

safetensors==0.4.2

sentencepiece==0.1.99

setuptools==68.2.2

six==1.16.0

soupsieve==2.5

sympy==1.12

tensorboard==2.14.0

tensorboard-data-server==0.7.2

tokenizers==0.13.3

torch==2.0.0

tqdm==4.66.2

transformers==4.34.0

triton==2.0.0

typing_extensions==4.11.0

tzdata==2024.1

urllib3==2.2.1

virtualenv==20.25.1

Werkzeug==3.0.2

wheel==0.41.2

wrapt==1.16.0

xxhash==3.4.1

yarl==1.9.4

\n[SEP]\n#finetune. py \n

import argparse

import os

import resource

from contextlib import nullcontext

from functools import partial

from typing import Optional, Tuple

import torch

import torch. distributed as dist

import torch. nn as nn

from attn import replace_with_flash_attention

# xuan ========================================

import sys; sys. path. append("..")

# ========================================

# datasets 2.18.0

# fsspec 2024.2.0

from data_utils import load_json, prepare_dataloader, save_json

from datasets import load_dataset, load_from_disk # , save_to_disk

from torch. optim import Optimizer

from torch. optim. lr_scheduler import _LRScheduler

from torch. utils. tensorboard import SummaryWriter

from tqdm import tqdm

from transformers. models. llama. configuration_llama import LlamaConfig

from transformers. models. llama. modeling_llama import LlamaForCausalLM

from transformers. models. llama. tokenization_llama import LlamaTokenizer

import colossalai

from colossalai. accelerator import get_accelerator

from colossalai. booster import Booster

from colossalai. booster. plugin import GeminiPlugin, HybridParallelPlugin, LowLevelZeroPlugin

from colossalai. cluster import DistCoordinator

from colossalai. lazy import LazyInitContext

from colossalai. nn. lr_scheduler import CosineAnnealingWarmupLR

from colossalai. nn. optimizer import HybridAdam

MODEL_CONFIGS = {

"7b": LlamaConfig(max_position_embeddings=4096),

"13b": LlamaConfig(

hidden_size=5120,

intermediate_size=13824,

num_hidden_layers=40,

num_attention_heads=40,

max_position_embeddings=4096,

"70b": LlamaConfig(

hidden_size=8192,

intermediate_size=28672,

num_hidden_layers=80,

num_attention_heads=64,

max_position_embeddings=4096,

num_key_value_heads=8,

}

def get_model_numel(model: nn. Module ) -> int:

return sum(p. numel () for p in model. parameters ())

def format_numel_str(numel: int) -> str:

B = 1024 ** 3

M = 1024 ** 2

K = 1024

if numel >= B:

return f"{numel / B:.2f} B"

elif numel >= M:

return f"{numel / M:.2f} M"

elif numel >= K:

return f"{numel / K:.2f} K"

else:

return f"{numel}"

def tokenize_batch_for_pretrain(batch, tokenizer: Optional[LlamaTokenizer] = None, max_length: int = 2048):

texts = [sample["text"] for sample in batch]

data = tokenizer(texts, return_tensors="pt", padding="max_length", truncation=True, max_length=max_length)

data = {k: v. cuda () for k, v in data. items ()}

data["labels"] = data["input_ids"].clone()

return data

def all_reduce_mean(tensor: torch. Tensor ) -> torch. Tensor :

dist. all_reduce (tensor, op=dist. ReduceOp. SUM)

tensor = tensor. data

tensor. div_ (dist. get_world_size ())

return tensor

def save(

booster: Booster,

model: nn. Module ,

optimizer: Optimizer,

lr_scheduler: _LRScheduler,

epoch: int,

step: int,

batch_size: int,

coordinator: DistCoordinator,

save_dir: str,

save_dir = os. path. join(save_dir, f"epoch{epoch}-step{step}")

os. makedirs (os. path. join(save_dir, "model"), exist_ok=True)

booster. save_model (model, os. path. join(save_dir, "model"), shard=True)

booster. save_optimizer (optimizer, os. path. join(save_dir, "optimizer"), shard=True)

booster. save_lr_scheduler (lr_scheduler, os. path. join(save_dir, "lr_scheduler"))

running_states = {

"epoch": epoch,

"step": step,

"sample_start_index": step * batch_size,

}

if coordinator. is_master ():

save_json(running_states, os. path. join(save_dir, "running_states. json "))

def load(

booster: Booster, model: nn. Module , optimizer: Optimizer, lr_scheduler: _LRScheduler, load_dir: str

) -> Tuple[int, int, int]:

booster. load_model (model, os. path. join(load_dir, "model"))

# booster. load_optimizer (optimizer, os. path. join(load_dir, "optimizer"))

# booster. load_lr_scheduler (lr_scheduler, os. path. join(load_dir, "lr_scheduler"))

running_states = load_json(os. path. join(load_dir, "running_states. json "))

return running_states["epoch"], running_states["step"], running_states["sample_start_index"]

def _criterion(outputs, inputs):

return outputs. loss

def encode_with_prompt_completion_format(example, tokenizer, max_seq_length):

'''

Here we assume each example has 'prompt' and 'completion' fields. We concatenate prompt and completion and tokenize them together because otherwise prompt will be padded/trancated

and it doesn't make sense to follow directly with the completion. '''

# if prompt doesn't end with space and completion doesn't start with space, add space

if not example['prompt'].endswith((' ', '\n', '\t')) and not example['completion'].startswith((' ', '\n', '\t')):

example_text = example['prompt'] + ' ' + example['completion']

else:

example_text = example['prompt'] + example['completion']

example_text = example_text + tokenizer. eos_token

tokenized_example = tokenizer(example_text, return_tensors='pt', max_length=max_seq_length, truncation=True)

input_ids = tokenized_example. input_ids

labels = input_ids. clone ()

tokenized_prompt = tokenizer(example['prompt'], return_tensors='pt', max_length=max_seq_length, truncation=True)

# mask the prompt part for avoiding loss

labels[:, :tokenized_prompt. input_ids. shape[1]] = -100

attention_mask = torch. ones_like (input_ids)

return {

'input_ids': input_ids. flatten (),

'labels': labels. flatten (),

'attention_mask': attention_mask. flatten (),

}

def encode_with_messages_format(example, tokenizer, max_seq_length):

'''

Here we assume each example has a 'messages' field Each message is a dict with 'role' and 'content' fields. We concatenate all messages with the roles as delimiters and tokenize them together. '''

messages = example['messages']

if len(messages) == 0:

raise ValueError('messages field is empty.')

def _concat_messages(messages):

message_text = ""

for message in messages:

if message["role"] == "system":

message_text += "<|system|>\n" + message["content"].strip() + "\n"

elif message["role"] == "user":

message_text += "<|user|>\n" + message["content"].strip() + "\n"

elif message["role"] == "assistant":

message_text += "<|assistant|>\n" + message["content"].strip() + tokenizer. eos_token + "\n"

else:

raise ValueError("Invalid role: {}".format(message["role"]))

return message_text

example_text = _concat_messages(messages).strip()

tokenized_example = tokenizer(example_text, return_tensors='pt', max_length=max_seq_length, truncation=True)

input_ids = tokenized_example. input_ids

labels = input_ids. clone ()

# mask the non-assistant part for avoiding loss

for message_idx, message in enumerate(messages):

if message["role"] != "assistant":

if message_idx == 0:

message_start_idx = 0

else:

message_start_idx = tokenizer(

_concat_messages(messages[:message_idx]), return_tensors='pt', max_length=max_seq_length,

truncation=True

).input_ids. shape [1]

if message_idx < len(messages) - 1 and messages[message_idx + 1]["role"] == "assistant":

# here we also ignore the role of the assistant

messages_so_far = _concat_messages(messages[:message_idx + 1]) + "<|assistant|>\n"

else:

messages_so_far = _concat_messages(messages[:message_idx + 1])

message_end_idx = tokenizer(

messages_so_far,

return_tensors='pt',

max_length=max_seq_length,

truncation=True

).input_ids. shape [1]

labels[:, message_start_idx:message_end_idx] = -100

if message_end_idx >= max_seq_length:

break

attention_mask = torch. ones_like (input_ids)

return {

'input_ids': input_ids. flatten (),

'labels': labels. flatten (),

'attention_mask': attention_mask. flatten (),

}

def main():

# ==============================

# Parse Arguments

# ==============================

parser = argparse. ArgumentParser ()

parser. add_argument ("-c", "--config", type=str, default="7b", help="Model configuration")

parser. add_argument (

"-p",

"--plugin",

choices=["gemini", "gemini_auto", "zero2", "zero2_cpu", "hybrid_parallel"],

default="gemini",

help="Choose which plugin to use",

)

parser. add_argument ("-e", "--num_epochs", type=int, default=1, help="Number of epochs")

parser. add_argument ("-b", "--batch_size", type=int, default=2, help="Local batch size")

parser. add_argument ("--lr", type=float, default=3e-4, help="Learning rate")

parser. add_argument ("--data_path", type=str, default="workspace/datasets/tulu_v2.jsonl ", help="dataset path")

parser. add_argument ("-w", "--weigth_decay", type=float, default=0.1 , help="Weight decay")

parser. add_argument ("-s", "--warmup_steps", type=int, default=2000, help="Warmup steps")

parser. add_argument ("-g", "--grad_checkpoint", action="store_true", help="Use gradient checkpointing")

parser. add_argument ("-l", "--max_length", type=int, default=4096, help="Max sequence length")

parser. add_argument ("-x", "--mixed_precision", default="fp16", choices=["fp16", "bf16"], help="Mixed precision")

parser. add_argument ("-i", "--save_interval", type=int, default=1000, help="Save interval")

parser. add_argument ("-o", "--save_dir", type=str, default="checkpoint", help="Checkpoint directory")

parser. add_argument ("-f", "--load", type=str, default=None, help="Load checkpoint")

parser. add_argument ("--grad_clip", type=float, default=1.0 , help="Gradient clipping")

parser. add_argument ("-t", "--tensorboard_dir", type=str, default="tb_logs", help="Tensorboard directory")

parser. add_argument ("-a", "--flash_attention", action="store_true", help="Use Flash Attention")

# xuan ================================================================================

parser. add_argument ("--accumulation_steps", default=16, help="accumulation steps")

parser. add_argument ("--expanded_model", default="", help="model path")

parser. add_argument ("--tokenizer_path", default="", help="model path")

# ================================================================================

args = parser. parse_args ()

# ==============================

# Initialize Distributed Training

# ==============================

colossalai. launch_from_torch ({})

coordinator = DistCoordinator()

# ==============================

# Initialize Booster

# ==============================

if args. plugin == "gemini":

plugin = GeminiPlugin(precision=args. mixed_precision , initial_scale=2 ** 16, max_norm=args. grad_clip )

elif args. plugin == "gemini_auto":

plugin = GeminiPlugin(

precision=args. mixed_precision , placement_policy="auto", initial_scale=2 ** 16, max_norm=args. grad_clip

)

elif args. plugin == "zero2":

plugin = LowLevelZeroPlugin(

stage=2, precision=args. mixed_precision , initial_scale=2 ** 16, max_norm=args. grad_clip

)

elif args. plugin == "zero2_cpu":

plugin = LowLevelZeroPlugin(

stage=2, precision=args. mixed_precision , initial_scale=2 ** 16, cpu_offload=True, max_norm=args. grad_clip

)

elif args. plugin == "hybrid_parallel":

plugin = HybridParallelPlugin(

tp_size=4,

pp_size=2,

num_microbatches=None,

microbatch_size=1,

enable_jit_fused=False,

zero_stage=0,

precision=args. mixed_precision ,

initial_scale=1,

)

else:

raise ValueError(f"Unknown plugin {args. plugin }")

booster = Booster(plugin=plugin)

use_pipeline = isinstance(booster. plugin , HybridParallelPlugin) and booster. plugin. pp_size > 1

is_pp_last_stage = use_pipeline and booster. plugin. stage_manager. is_last_stage ()

print_flag = (not use_pipeline and coordinator. is_master ()) or (use_pipeline and is_pp_last_stage)

# ==============================

# Initialize Tensorboard

# ==============================

if print_flag:

os. makedirs (args. tensorboard_dir , exist_ok=True)

writer = SummaryWriter(args. tensorboard_dir )

# ==============================

# Initialize Tokenizer, Dataset and Dataloader

# ==============================

from transformers import AutoTokenizer

tokenizer = AutoTokenizer. from_pretrained (args. tokenizer_path , use_fast=False)

tokenizer. pad_token = tokenizer. unk_token

# ================================================================

train_file = args. data_path

data_files = {}

dataset_args = {}

if train_file is not None:

data_files["train"] = train_file

raw_datasets = load_dataset(

"json",

data_files=data_files,

**dataset_args,

)

# Preprocessing the datasets. if "prompt" in raw_datasets["train"].column_names and "completion" in raw_datasets["train"].column_names:

encode_function = partial(

encode_with_prompt_completion_format,

tokenizer=tokenizer,

max_seq_length=args. max_length ,

)

elif "messages" in raw_datasets["train"].column_names:

encode_function = partial(

encode_with_messages_format,

tokenizer=tokenizer,

max_seq_length=args. max_length ,

)

else:

raise ValueError("You need to have either 'prompt'&'completion' or 'messages' in your column names.")

# with accelerator. main_process_first ():

# if coordinator. is_master ():

lm_datasets = raw_datasets. map (

encode_function,

batched=False,

num_proc=16,

load_from_cache_file=True,

remove_columns=[name for name in raw_datasets["train"].column_names if

name not in ["input_ids", "labels", "attention_mask"]],

desc="Tokenizing and reformatting instruction data",

)

lm_datasets. set_format (type="pt")

lm_datasets = lm_datasets. filter (lambda example: (example['labels'] != -100).any())

dist. barrier ()

train_dataset = lm_datasets["train"]

train_ds = train_dataset

# ==============================

# Initialize Model, Optimizer and LR Scheduler

# ==============================

config = MODEL_CONFIGS[args. config ]

init_ctx = (

LazyInitContext(default_device=get_accelerator().get_current_device())

if isinstance(plugin, GeminiPlugin)

else nullcontext()

)

with init_ctx:

from transformers import AutoModelForCausalLM, AutoConfig

model = AutoModelForCausalLM. from_pretrained (args. expanded_model , torch_dtype=torch. float16 )#, device_map='cpu')

for name, weight in model. named_parameters ():

weight. requires_grad = True

from transformers import DataCollatorForSeq2Seq

dataloader = prepare_dataloader(

lm_datasets["train"],

batch_size=args. batch_size ,

shuffle=True,

drop_last=True,

collate_fn=DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model, padding="longest"),

)

# =======================================================================

if args. grad_checkpoint :

model. gradient_checkpointing_enable ()

model_numel = get_model_numel(model)

coordinator. print_on_master (f"Model params: {format_numel_str(model_numel)}")

optimizer = HybridAdam(model. parameters (), lr=args. lr , betas=(0.9 , 0.95 ), weight_decay=args. weigth_decay )

lr_scheduler = CosineAnnealingWarmupLR(

optimizer, total_steps=args. num_epochs * len(dataloader), warmup_steps=args. warmup_steps , eta_min=0.1 * args. lr

)

default_dtype = torch. float16 if args. mixed_precision == "fp16" else torch. bfloat16

torch. set_default_dtype (default_dtype)

model, optimizer, _, dataloader, lr_scheduler = booster. boost (

model, optimizer, dataloader=dataloader, lr_scheduler=lr_scheduler

)

torch. set_default_dtype (torch. float )

coordinator. print_on_master (f"Booster init max CUDA memory: {torch. cuda. max_memory_allocated() / 1024 ** 2:.2f} MB")

coordinator. print_on_master (

f"Booster init max CPU memory: {resource. getrusage (resource. RUSAGE_SELF ).ru_maxrss / 1024:.2f} MB"

)

# load checkpoint if specified

start_epoch = 0

start_step = 0

sampler_start_idx = 0

if args. load is not None:

coordinator. print_on_master ("Loading checkpoint")

start_epoch, start_step, sampler_start_idx = load(booster, model, optimizer, lr_scheduler, args. load )

coordinator. print_on_master (f"Loaded checkpoint {args. load } at epoch {start_epoch} step {start_step}")

num_steps_per_epoch = len(dataloader)

# if resume training, set the sampler start index to the correct value

dataloader. sampler. set_start_index(sampler_start_idx)

for epoch in range(start_epoch, args. num_epochs ):

dataloader. sampler. set_epoch(epoch)

dataloader_iter = iter(dataloader)

for step in range(start_step, num_steps_per_epoch):

if use_pipeline:

outputs = booster. execute_pipeline (dataloader_iter, model, _criterion, optimizer, return_loss=True)

loss = outputs["loss"]

else:

batch = next(dataloader_iter)

batch = batch. to (get_accelerator().get_current_device())

outputs = model(**batch)

loss = outputs[0]

booster. backward (loss, optimizer)

if (step + 1) % args. accumulation_steps == 0:

optimizer. step () # Update parameters

lr_scheduler. step ()

optimizer. zero_grad () # Reset gradients

if step + 1 == num_steps_per_epoch - 1:

optimizer. step ()

lr_scheduler. step ()

optimizer. zero_grad ()

if not use_pipeline:

all_reduce_mean(loss)

if print_flag:

writer. add_scalar ("loss", loss. item (), epoch * num_steps_per_epoch + step)

print("Epoch: {}, step: {}, loss: {:.3f}".format(

epoch,

epoch * num_steps_per_epoch + step,

loss. item ()

))

if args. save_interval > 0 and (step + 1) % args. save_interval == 0:

coordinator. print_on_master (f"Saving checkpoint")

save(

booster,

model,

optimizer,

lr_scheduler,

epoch,

step + 1,

args. batch_size ,

coordinator,

args. save_dir ,

)

coordinator. print_on_master (f"Saved checkpoint at epoch {epoch} step {step + 1}")

# the continue epochs are not resumed, so we need to reset the sampler start index and start step

dataloader. sampler. set_start_index(0)

start_step = 0

coordinator. print_on_master (f"Max CUDA memory usage: {torch. cuda. max_memory_allocated() / 1024 ** 2:.2f} MB")

if __name__ == "__main__":

main()

\n[SEP]\n#inference. py \n

import torch

from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

model_name = './model'

tokenizer = AutoTokenizer. from_pretrained (model_name)

model = AutoModelForCausalLM. from_pretrained (

model_name,

torch_dtype=torch. bfloat16 ,

device_map="auto",

trust_remote_code=True,

)

pipe = pipeline(

"text-generation",

model=model,

tokenizer = tokenizer,

torch_dtype=torch. bfloat16 ,

device_map="auto"

)

prompt = "Can you explain the concept of regularization in machine learning?"

sequences = pipe(

prompt,

do_sample=True,

max_new_tokens=100,

temperature=0.7 ,

top_k=50,

top_p=0.95 ,

num_return_sequences=1,

)

print(sequences[0]['generated_text'])

\n[SEP]\n#attn. py \n

import math

from types import MethodType

from typing import Optional, Tuple

import torch

import torch. nn as nn

import torch. nn. functional as F

from einops import rearrange

from transformers. models. llama. configuration_llama import LlamaConfig

from transformers. models. llama. modeling_llama import (

LlamaAttention,

LlamaForCausalLM,

LlamaModel,

LlamaRMSNorm,

apply_rotary_pos_emb,

repeat_kv,

)

from colossalai. accelerator import get_accelerator

from colossalai. logging import get_dist_logger

logger = get_dist_logger()

if get_accelerator().name == "cuda":

from flash_attn. bert_padding import pad_input, unpad_input

from flash_attn. flash_attn_interface import flash_attn_func, flash_attn_varlen_kvpacked_func

from flash_attn. ops. rms_norm import rms_norm

def _prepare_decoder_attention_mask(

self: LlamaModel,

attention_mask: torch. BoolTensor ,

input_shape: torch. Size ,

inputs_embeds: torch. Tensor ,

past_key_values_length: int,

) -> Optional[torch. Tensor ]:

"""

Decoder attetion mask

"""

if past_key_values_length > 0 and attention_mask is not None:

attention_mask = torch. cat (

tensors=(

torch. full (

size=(input_shape[0], past_key_values_length),

fill_value=True,

dtype=attention_mask. dtype ,

device=attention_mask. device ,

attention_mask,

dim=-1,

) # (bsz, past_key_values_length + q_len)

if attention_mask is not None and torch. all (attention_mask):

return None # Faster

return attention_mask

def attention_forward(

self: LlamaAttention,

hidden_states: torch. Tensor ,

attention_mask: Optional[torch. Tensor ] = None,

position_ids: Optional[torch. LongTensor ] = None,

past_key_value: Optional[Tuple[torch. Tensor ]] = None,

output_attentions: bool = False,

use_cache: bool = False,

**kwargs,

) -> Tuple[torch. Tensor , Optional[torch. Tensor ], Optional[Tuple[torch. Tensor ]]]:

"""

Re-define LLaMA-2 `LlamaAttention` forward method using flash-attention. """

if output_attentions:

logger. warning (

"Argument `output_attentions` is not supported for flash-attention patched `LlamaAttention`, "

"return `None` instead."

)

bsz, q_len, _ = hidden_states. size ()

if self. config. pretraining_tp > 1:

q_slicing, kv_slicing = (

dim // self. config. pretraining_tp

for dim in (

self. num_heads * self. head_dim ,

self. num_key_value_heads * self. head_dim ,

)

) # `Tuple[int, int]`

q_slices, k_slices, v_slices = (

proj. weight. split(slicing, dim=0)

for proj, slicing in (

(self. q_proj , q_slicing),

(self. k_proj , kv_slicing),

(self. v_proj , kv_slicing),

)

) # Tuple[Tuple[torch. Tensor ], Tuple[torch. Tensor ], Tuple[torch. Tensor ]]

q, k, v = (

torch. cat (

[F. linear (hidden_states, slices[i]) for i in range(self. config. pretraining_tp)],

dim=-1,

)

for slices in (q_slices, k_slices, v_slices)

)

# `Tuple[torch. Tensor , torch. Tensor , torch. Tensor ]` of shape:

# (bsz, q_len, num_heads * head_dim),

# (bsz, q_len, num_key_value_heads * head_dim),

# (bsz, q_len, num_key_value_heads * head_dim)

else:

q, k, v = (proj(hidden_states) for proj in (self. q_proj , self. k_proj , self. v_proj ))

# `Tuple[torch. Tensor , torch. Tensor , torch. Tensor ]` of shape:

# (bsz, q_len, num_heads * head_dim),

# (bsz, q_len, num_key_value_heads * head_dim),

# (bsz, q_len, num_key_value_heads * head_dim)

# (bsz, q_len, num_heads * head_dim) -> (bsz, num_heads, q_len, head_dim);

# (bsz, q_len, num_key_value_heads * head_dim) -> (bsz, num_key_value_heads, q_len, head_dim);

# (bsz, q_len, num_key_value_heads * head_dim) -> (bsz, num_key_value_heads, q_len, head_dim)

q, k, v = (

states. view (bsz, q_len, num_heads, self. head_dim ).transpose(1, 2)

for states, num_heads in (

(q, self. num_heads ),

(k, self. num_key_value_heads ),

(v, self. num_key_value_heads ),

)

kv_len = k. shape [-2] # initially, `kv_len` == `q_len`

past_kv_len = 0

if past_key_value is not None:

# if `past_key_value` is not None, `kv_len` > `q_len`.

past_kv_len = past_key_value[0].shape[-2]

kv_len += past_kv_len

# two `torch. Tensor ` objs of shape (1, 1, kv_len, head_dim)

cos, sin = self. rotary_emb (v, seq_len=kv_len)

# (bsz, num_heads, q_len, head_dim), (bsz, num_key_value_heads, q_len, head_dim)

q, k = apply_rotary_pos_emb(q=q, k=k, cos=cos, sin=sin, position_ids=position_ids)

if past_key_value is not None:

# reuse k, v, self_attention

k = torch. cat ([past_key_value[0], k], dim=2)

v = torch. cat ([past_key_value[1], v], dim=2)

past_key_value = (k, v) if use_cache else None

# repeat k/v heads if n_kv_heads < n_heads

k = repeat_kv(hidden_states=k, n_rep=self. num_key_value_groups )

# (bsz, num_key_value_heads, q_len, head_dim) -> (bsz, num_heads, q_len, head_dim)

v = repeat_kv(hidden_states=v, n_rep=self. num_key_value_groups )

# (bsz, num_key_value_heads, q_len, head_dim) -> (bsz, num_heads, q_len, head_dim)

key_padding_mask = attention_mask

# (bsz, num_heads, q_len, head_dim) -> (bsz, q_len, num_heads, head_dim)

q, k, v = (states. transpose (1, 2) for states in (q, k, v))

if past_kv_len > 0:

q = torch. cat (

tensors=(

torch. full (

size=(bsz, past_kv_len, self. num_heads , self. head_dim ),

fill_value=0.0 ,

dtype=q. dtype ,

device=q. device ,

dim=1,

) # (bsz, past_kv_len + q_len, num_heads, head_dim)

if key_padding_mask is None:

# (bsz, past_kv_len + q_len, num_heads, head_dim)

output = flash_attn_func(q=q, k=k, v=v, dropout_p=0.0 , softmax_scale=None, causal=True) # (bsz, )

output = rearrange(

output, pattern="...

h d ->...

(h d)"

) # (bsz, past_kv_len + q_len, num_heads * head_dim)

else:

q, indices, cu_q_lens, max_q_len = unpad_input(hidden_states=q, attention_mask=key_padding_mask)

kv, _, cu_kv_lens, max_kv_len = unpad_input(

hidden_states=torch. stack (tensors=(k, v), dim=2),

attention_mask=key_padding_mask,

)

output_unpad = flash_attn_varlen_kvpacked_func(

q=q,

kv=kv,

cu_seqlens_q=cu_q_lens,

cu_seqlens_k=cu_kv_lens,

max_seqlen_q=max_q_len,

max_seqlen_k=max_kv_len,

dropout_p=0.0 ,

softmax_scale=None,

causal=True,

)

output = pad_input(

hidden_states=rearrange(output_unpad, pattern="nnz h d -> nnz (h d)"),

indices=indices,

batch=bsz,

seqlen=past_kv_len + q_len,

) # (bsz, past_kv_len + q_len, num_heads * head_dim)

if past_kv_len > 0:

# Strip off the zero query outputs. output = output[:, past_kv_len:,...] # (bsz, q_len, num_heads * head_dim)

output = self. o_proj (output) # (bsz, q_len, hidden_size)

return output, None, past_key_value

def rms_norm_forward(self: LlamaRMSNorm, hidden_states: torch. Tensor ) -> torch. Tensor :

"""

Formard function for RMS Norm

"""

return rms_norm(x=hidden_states, weight=self. weight , epsilon=self. variance_epsilon )

def replace_with_flash_attention(model: LlamaForCausalLM) -> None:

for name, module in model. named_modules ():

if isinstance(module, LlamaAttention):

module. forward = MethodType(attention_forward, module)

if isinstance(module, LlamaModel):

module._prepare_decoder_attention_mask = MethodType(_prepare_decoder_attention_mask, module)

if isinstance(module, LlamaRMSNorm):

module. forward = MethodType(rms_norm_forward, module)

elif get_accelerator().name == "npu":

import torch_npu

class NPULlamaAttention(LlamaAttention):

use_flash: bool = True

def __init__(self, config: LlamaConfig):

super().__init__(config)

self. setup ()

def setup(self):

self._softmax_scale = 1 / math. sqrt (self. head_dim )

def forward(

self,

hidden_states: torch. Tensor ,

attention_mask: Optional[torch. Tensor ] = None,

position_ids: Optional[torch. LongTensor ] = None,

past_key_value: Optional[Tuple[torch. Tensor ]] = None,

output_attentions: bool = False,

use_cache: bool = False,

) -> Tuple[torch. Tensor , Optional[torch. Tensor ], Optional[Tuple[torch. Tensor ]]]:

bsz, q_len, _ = hidden_states. size ()

if self. config. pretraining_tp > 1:

key_value_slicing = (self. num_key_value_heads * self. head_dim ) // self. config. pretraining_tp

query_slices = self. q_proj. weight. split (

(self. num_heads * self. head_dim ) // self. config. pretraining_tp, dim=0

)

key_slices = self. k_proj. weight. split (key_value_slicing, dim=0)

value_slices = self. v_proj. weight. split (key_value_slicing, dim=0)

query_states = [F. linear (hidden_states, query_slices[i]) for i in range(self. config. pretraining_tp)]

query_states = torch. cat (query_states, dim=-1)

key_states = [F. linear (hidden_states, key_slices[i]) for i in range(self. config. pretraining_tp)]

key_states = torch. cat (key_states, dim=-1)

value_states = [F. linear (hidden_states, value_slices[i]) for i in range(self. config. pretraining_tp)]

value_states = torch. cat (value_states, dim=-1)

else:

query_states = self. q_proj (hidden_states)

key_states = self. k_proj (hidden_states)

value_states = self. v_proj (hidden_states)

query_states = query_states. view (bsz, q_len, self. num_heads , self. head_dim ).transpose(1, 2)

key_states = key_states. view (bsz, q_len, self. num_key_value_heads , self. head_dim ).transpose(1, 2)

value_states = value_states. view (bsz, q_len, self. num_key_value_heads , self. head_dim ).transpose(1, 2)

kv_seq_len = key_states. shape [-2]

if past_key_value is not None:

kv_seq_len += past_key_value[0].shape[-2]

cos, sin = self. rotary_emb (value_states, seq_len=kv_seq_len)

query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)

if past_key_value is not None:

# reuse k, v, self_attention

key_states = torch. cat ([past_key_value[0], key_states], dim=2)

value_states = torch. cat ([past_key_value[1], value_states], dim=2)

past_key_value = (key_states, value_states) if use_cache else None

key_states = repeat_kv(key_states, self. num_key_value_groups )

value_states = repeat_kv(value_states, self. num_key_value_groups )

if not self. use_flash :

attn_weights = torch. matmul (query_states, key_states. transpose (2, 3)) / math. sqrt (self. head_dim )

if attn_weights. size () != (bsz, self. num_heads , q_len, kv_seq_len):

raise ValueError(

f"Attention weights should be of size {(bsz, self. num_heads , q_len, kv_seq_len)}, but is"

f" {attn_weights. size ()}"

)

if attention_mask is not None:

if attention_mask. size () != (bsz, 1, q_len, kv_seq_len):

raise ValueError(

f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask. size ()}"

)

attn_weights = attn_weights + attention_mask

# upcast attention to fp32

attn_weights = nn. functional. softmax(attn_weights, dim=-1, dtype=torch. float32 ).to(query_states. dtype )

attn_output = torch. matmul (attn_weights, value_states)

else:

attn_output, *_ = torch_npu. npu_fusion_attention (

query_states,

key_states,

value_states,

self. num_heads ,

"BNSD",

atten_mask=attention_mask. bool (),

scale=self._softmax_scale ,

padding_mask=None,

pre_tockens=65535,

next_tockens=0,

keep_prob=1.0 ,

inner_precise=0,

)

if attn_output. size () != (bsz, self. num_heads , q_len, self. head_dim ):

raise ValueError(

f"`attn_output` should be of size {(bsz, self. num_heads , q_len, self. head_dim )}, but is"

f" {attn_output. size ()}"

)

attn_output = attn_output. transpose (1, 2).contiguous()

attn_output = attn_output. reshape (bsz, q_len, self. hidden_size )

if self. config. pretraining_tp > 1:

attn_output = attn_output. split (self. hidden_size // self. config. pretraining_tp, dim=2)

o_proj_slices = self. o_proj. weight. split (self. hidden_size // self. config. pretraining_tp, dim=1)

attn_output = sum(

[F. linear (attn_output[i], o_proj_slices[i]) for i in range(self. config. pretraining_tp)]

)

else:

attn_output = self. o_proj (attn_output)

if not output_attentions:

attn_weights = None

return attn_output, attn_weights, past_key_value

class NPURMSNorm(LlamaRMSNorm):

def forward(self, hidden_states):

return torch_npu. npu_rms_norm (hidden_states, self. weight , epsilon=self. variance_epsilon )[0]

def replace_with_flash_attention(model: LlamaForCausalLM) -> None:

for name, module in model. named_modules ():

if isinstance(module, LlamaAttention):

module.__class__ = NPULlamaAttention

module. setup ()

if isinstance(module, LlamaRMSNorm):

module.__class__ = NPURMSNorm

\n[SEP]\n#pretrain. py \n

import argparse

import os

import resource

from contextlib import nullcontext

from functools import partial

from typing import Optional, Tuple

import torch

import torch. distributed as dist

import torch. nn as nn

from attn import replace_with_flash_attention

import sys

sys. path. append("..")

# datasets 2.18.0

# fsspec 2024.2.0

from data_utils import load_json, prepare_dataloader, save_json

from datasets import load_dataset, load_from_disk

from torch. optim import Optimizer

from torch. optim. lr_scheduler import _LRScheduler

from torch. utils. tensorboard import SummaryWriter

from tqdm import tqdm

from transformers. models. llama. configuration_llama import LlamaConfig

from transformers. models. llama. modeling_llama import LlamaForCausalLM

from transformers. models. llama. tokenization_llama import LlamaTokenizer

import colossalai

from colossalai. accelerator import get_accelerator

from colossalai. booster import Booster

from colossalai. booster. plugin import GeminiPlugin, HybridParallelPlugin, LowLevelZeroPlugin

from colossalai. cluster import DistCoordinator

from colossalai. lazy import LazyInitContext

from colossalai. nn. lr_scheduler import CosineAnnealingWarmupLR

from colossalai. nn. optimizer import HybridAdam

MODEL_CONFIGS = {

"7b": LlamaConfig(max_position_embeddings=4096),

"13b": LlamaConfig(

hidden_size=5120,

intermediate_size=13824,

num_hidden_layers=40,

num_attention_heads=40,

max_position_embeddings=4096,

"70b": LlamaConfig(

hidden_size=8192,

intermediate_size=28672,

num_hidden_layers=80,

num_attention_heads=64,

max_position_embeddings=4096,

num_key_value_heads=8,

}

def get_model_numel(model: nn. Module ) -> int:

return sum(p. numel () for p in model. parameters ())

def format_numel_str(numel: int) -> str:

B = 1024**3

M = 1024**2

K = 1024

if numel >= B:

return f"{numel / B:.2f} B"

elif numel >= M:

return f"{numel / M:.2f} M"

elif numel >= K:

return f"{numel / K:.2f} K"

else:

return f"{numel}"

def tokenize_batch_for_pretrain(batch, tokenizer: Optional[LlamaTokenizer] = None, max_length: int = 2048):

texts = [sample["text"] for sample in batch]

data = tokenizer(texts, return_tensors="pt", padding="max_length", truncation=True, max_length=max_length)

data = {k: v. cuda () for k, v in data. items ()}

data["labels"] = data["input_ids"].clone()

return data

def all_reduce_mean(tensor: torch. Tensor ) -> torch. Tensor :

dist. all_reduce (tensor, op=dist. ReduceOp. SUM)

tensor = tensor. data

tensor. div_ (dist. get_world_size ())

return tensor

def save(

booster: Booster,

model: nn. Module ,

optimizer: Optimizer,

lr_scheduler: _LRScheduler,

epoch: int,

step: int,

batch_size: int,

coordinator: DistCoordinator,

save_dir: str,

save_dir = os. path. join(save_dir, f"epoch{epoch}-step{step}")

os. makedirs (os. path. join(save_dir, "model"), exist_ok=True)

booster. save_model (model, os. path. join(save_dir, "model"), shard=True)

booster. save_optimizer (optimizer, os. path. join(save_dir, "optimizer"), shard=True)

booster. save_lr_scheduler (lr_scheduler, os. path. join(save_dir, "lr_scheduler"))

running_states = {

"epoch": epoch,

"step": step,

"sample_start_index": step * batch_size,

}

if coordinator. is_master ():

save_json(running_states, os. path. join(save_dir, "running_states. json "))

def load(

booster: Booster, model: nn. Module , optimizer: Optimizer, lr_scheduler: _LRScheduler, load_dir: str

) -> Tuple[int, int, int]:

booster. load_model (model, os. path. join(load_dir, "model"))

running_states = load_json(os. path. join(load_dir, "running_states. json "))

return running_states["epoch"], running_states["step"], running_states["sample_start_index"]

def _criterion(outputs, inputs):

return outputs. loss

def main():

# ==============================

# Parse Arguments

# ==============================

parser = argparse. ArgumentParser ()

parser. add_argument ("-c", "--config", type=str, default="7b", help="Model configuration")

parser. add_argument (

"-p",

"--plugin",

choices=["gemini", "gemini_auto", "zero2", "zero2_cpu", "hybrid_parallel"],

default="gemini",

help="Choose which plugin to use",

)

parser. add_argument (

"-d", "--dataset", type=str, default="togethercomputer/RedPajama-Data-1T-Sample", help="Data set path"

)

parser. add_argument ("--cache_path", type=str, default="workspace/.cache/huggingface/datasets", help="cache path")