Initial Commit

.gitignore

@@ -0,0 +1,2 @@
+lightning_logs
+.env

.vscode/launch.json

@@ -0,0 +1,28 @@
+{
+    // 使用 IntelliSense 了解相关属性。 
+    // 悬停以查看现有属性的描述。
+    // 欲了解更多信息，请访问: https://go.microsoft.com/fwlink/?linkid=830387
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "Python: train",
+            "type": "python",
+            "request": "launch",
+            "program": "train.py",
+            "args": [
+                "--dataset_name", "wikitext:wikitext-2-v1",
+            ],
+            "console": "integratedTerminal",
+            "justMyCode": true
+        },
+        {
+            "name": "Python: generate",
+            "type": "python",
+            "request": "launch",
+            "program": "generate.py",
+            "args": [],
+            "console": "integratedTerminal",
+            "justMyCode": true
+        }
+    ]
+}

generate.py

@@ -0,0 +1,94 @@
+import argparse
+import os
+from typing import List, Union
+
+import torch
+from transformers import (
+    AutoModel,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    PreTrainedModel,
+    PreTrainedTokenizer,
+)
+
+
+def load_model(model_name_or_path: Union[str, os.PathLike]) -> PreTrainedModel:
+    try:
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name_or_path, trust_remote_code=True
+        )
+    except ValueError:
+        model = AutoModel.from_pretrained(model_name_or_path, trust_remote_code=True)
+    return model
+
+
+def load_tokenizer(model_name_or_path: Union[str, os.PathLike]) -> PreTrainedTokenizer:
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_name_or_path, padding_side='left', trust_remote_code=True
+    )
+    if tokenizer.pad_token_id is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    return tokenizer
+
+
+def eval_prompts(
+    model: PreTrainedModel, tokenizer: PreTrainedTokenizer, prompts: List[str]
+) -> List[str]:
+    inputs = tokenizer(
+        prompts, padding=True, return_tensors='pt', return_attention_mask=True
+    )
+    inputs['position_ids'] = inputs.attention_mask.cumsum(-1) - 1
+    inputs['position_ids'].masked_fill_(inputs.attention_mask == 0, 1)
+    inputs['attention_mask'] = (
+        inputs.attention_mask.unsqueeze(1) * inputs.attention_mask.unsqueeze(2)
+    ).tril()
+    inputs = inputs.to(model.device)
+    with torch.inference_mode():
+        output_ids = model.generate(
+            **inputs,
+            do_sample=False,
+            num_beams=16,
+            max_new_tokens=100,
+            eos_token_id=tokenizer.eos_token_id,
+            early_stopping=True,
+        )
+    completes = tokenizer.batch_decode(
+        output_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
+    )
+    return completes
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model_name_or_path",
+        type=str,
+        help="Name of or path to model",
+        default='gpt2',
+    )
+    args = parser.parse_args()
+    return args
+
+
+if __name__ == '__main__':
+    args = parse_args()
+
+    device = torch.device(0)
+
+    model = load_model(args.model_name_or_path)
+    tokenizer = load_tokenizer(args.model_name_or_path)
+
+    model = model.to(device)
+    prompts = [
+        "Shall I compare thee to a summer's day? Thou art more lovely and more temperate.",
+        "Shall I compare thee to a summer's day? Thou art more lovely and",
+        "Belle! C'est un mot qu'on dirait inventé pour elle.",
+        "Belle! C'est un mot qu'on dirait inventé",
+        "这是一个最好的时代，这是一个最坏的时代。",
+        "这是一个最好的时代，这是一个最坏的",
+    ]
+    completes = eval_prompts(model, tokenizer, prompts)
+
+    for prompt, complete in zip(prompts, completes):
+        print("[p]", prompt)
+        print("[c]", complete)

requirements.txt

@@ -0,0 +1,5 @@
+datasets==2.11.0
+deepspeed==0.9.1
+pytorch-lightning==2.0.2
+torch==2.0.0
+transformers==4.28.1

train.py

@@ -0,0 +1,288 @@
+import argparse
+import os
+from functools import cache, partial
+from itertools import chain
+from typing import Dict, Optional, Tuple, Union
+
+import datasets
+import pytorch_lightning as pl
+import torch
+import torchmetrics
+from pytorch_lightning import strategies
+from torch.utils.data import ConcatDataset, DataLoader
+from transformers import (
+    AutoConfig,
+    AutoModel,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    BatchEncoding,
+    DefaultDataCollator,
+    PreTrainedModel,
+    PreTrainedTokenizer,
+    set_seed,
+)
+
+
+def init_model(model_name: Union[str, os.PathLike]) -> PreTrainedModel:
+    config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
+    try:
+        model = AutoModelForCausalLM.from_config(config, trust_remote_code=True)
+    except ValueError:
+        model = AutoModel.from_config(config, trust_remote_code=True)
+    return model
+
+
+def load_tokenizer(
+    tokenizer_name_or_path: Union[str, os.PathLike]
+) -> PreTrainedTokenizer:
+    tokenizer = AutoTokenizer.from_pretrained(
+        tokenizer_name_or_path, padding_side='left', trust_remote_code=True
+    )
+    if tokenizer.pad_token_id is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    return tokenizer
+
+
+def split_raw_dataset(
+    raw_dataset: datasets.DatasetDict,
+) -> Tuple[datasets.Dataset, datasets.Dataset]:
+    if 'validation' in raw_dataset:
+        train_dataset, val_dataset = raw_dataset['train'], raw_dataset['validation']
+    else:
+        raw_dataset = raw_dataset['train'].train_test_split(
+            test_size=0.05, seed=args.seed
+        )
+        train_dataset, val_dataset = raw_dataset['train'], raw_dataset['test']
+    return train_dataset, val_dataset
+
+
+def process_dataset(
+    dataset: datasets.Dataset, tokenizer: PreTrainedTokenizer
+) -> datasets.Dataset:
+    def group_texts(examples: Dict[str, list], block_size: int = 512) -> BatchEncoding:
+        concatenated_examples = {k: list(chain(*examples[k])) for k in examples.keys()}
+        total_length = len(concatenated_examples[list(examples.keys())[0]])
+        total_length = (total_length // block_size) * block_size
+        result = {
+            k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
+            for k, t in concatenated_examples.items()
+        }
+        result['labels'] = result['input_ids'].copy()
+        result = BatchEncoding(result)
+        return result
+
+    def tokenize_inputs(
+        examples: Dict[str, list],
+        tokenizer: PreTrainedTokenizer,
+        column_name: str = 'text',
+    ) -> BatchEncoding:
+        return tokenizer(examples[column_name], return_attention_mask=False)
+
+    dataset_column_names = list(dataset.features)
+    dataset = dataset.map(
+        partial(
+            tokenize_inputs,
+            tokenizer=tokenizer,
+            column_name=dataset_column_names[0],
+        ),
+        batched=True,
+        num_proc=args.num_proc,
+        remove_columns=dataset_column_names,
+    ).map(
+        partial(group_texts, block_size=tokenizer.model_max_length),
+        batched=True,
+        num_proc=args.num_proc,
+    )
+    return dataset
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model_name",
+        type=str,
+        help="Name of or path to model",
+        default='gpt2',
+    )
+    parser.add_argument("--fp16", help="Enable fp16", action="store_true")
+    parser.add_argument("--bf16", help="Enable bf16", action="store_true")
+    parser.add_argument(
+        "--tokenizer_name_or_path",
+        type=str,
+        help="Name of or path to tokenizer",
+        default=None,
+    )
+    parser.add_argument(
+        "--dataset_name",
+        nargs='+',
+        type=str,
+        help="Name(s) of dataset. To specify a config, pass a <dataset_name>:<dataset_config_name>",
+        default=["wikitext:wikitext-2-v1"],
+    )
+    parser.add_argument(
+        "--train_batch_size",
+        type=int,
+        help="Batch size of training",
+        default=8,
+    )
+    parser.add_argument(
+        "--val_batch_size",
+        type=int,
+        help="Batch size of validating",
+        default=16,
+    )
+    parser.add_argument(
+        "--num_proc",
+        type=str,
+        help="Number of data processes",
+        default=16,
+    )
+    parser.add_argument(
+        "--resume_from_ckpt_path",
+        type=str,
+        help="Checkpoint file path to resume from",
+        default=None,
+    )
+    parser.add_argument(
+        "--seed",
+        type=str,
+        help="Random seed",
+        default=42,
+    )
+    args = parser.parse_args()
+    return args
+
+
+class LitModule(pl.LightningModule):
+    def __init__(self, model_name: str):
+        super().__init__()
+        self.save_hyperparameters()
+        self.llm = self.register_core_module(init_model(model_name))
+        self.metric_loss = torchmetrics.MeanMetric()
+        self.metric_accuracy = torchmetrics.Accuracy(
+            task='multiclass',
+            num_classes=self.llm.config.vocab_size,
+        )
+
+    @cache
+    def get_tril_matrix(
+        self, block_size: int, batch_size: Optional[int] = None
+    ) -> torch.Tensor:
+        matrix = torch.ones(block_size, block_size).tril()
+        if batch_size is not None:
+            matrix = matrix.repeat(batch_size, 1, 1)
+        return matrix
+
+    def register_core_module(self, module: torch.nn.Module) -> torch.nn.Module:
+        object.__setattr__(self, '__core_module__', module)
+        return module
+
+    def training_step(self, batch: Dict[str, torch.Tensor], batch_idx):
+        batch_size, block_size = batch['input_ids'].shape
+        batch['attention_mask'] = self.get_tril_matrix(
+            block_size, batch_size=batch_size
+        ).to(self.device)
+        outputs = self.llm(**batch, return_dict=True)
+        loss = outputs.loss
+
+        self.log('train_loss', loss, rank_zero_only=True)
+
+        return loss
+
+    def validation_step(self, batch: Dict[str, torch.Tensor], batch_idx):
+        outputs = self.llm(**batch, return_dict=True)
+        loss = outputs.loss
+        logits = outputs.logits[..., :-1, :]
+        labels = batch['labels'][..., 1:]
+
+        self.metric_loss.update(loss)
+
+        label_mask = labels != -100
+        self.metric_accuracy.update(logits[label_mask], labels[label_mask])
+
+    def on_validation_epoch_end(self) -> None:
+        self.log('val_loss', self.metric_loss, rank_zero_only=True)
+        self.log('accuracy', self.metric_accuracy, rank_zero_only=True)
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.AdamW(self.trainer.model.parameters(), lr=0.0001)
+        return optimizer
+
+    def configure_callbacks(self):
+        checkpoint_callback = pl.callbacks.ModelCheckpoint(
+            monitor='accuracy',
+            mode='max',
+            filename='{epoch:02d}-{accuracy:.4f}',
+        )
+        early_stop_callback = pl.callbacks.EarlyStopping(
+            monitor='accuracy',
+            min_delta=0.001,
+            patience=3,
+            mode='max',
+            stopping_threshold=1,
+        )
+        return [checkpoint_callback, early_stop_callback]
+
+
+if __name__ == '__main__':
+    args = parse_args()
+
+    if args.tokenizer_name_or_path is None:
+        args.tokenizer_name_or_path = args.model_name
+
+    set_seed(args.seed)
+
+    # lightning module
+    lit_module = LitModule(args.model_name)
+
+    # datasets
+    tokenizer = load_tokenizer(args.tokenizer_name_or_path)
+    train_dataset_list = []
+    val_dataset_list = []
+    for dataset_name in args.dataset_name:
+        dataset_args = dataset_name.split(':')
+        raw_dataset = datasets.load_dataset(*dataset_args)
+        train_dataset, val_dataset = split_raw_dataset(raw_dataset)
+        train_dataset = process_dataset(train_dataset, tokenizer)
+        val_dataset = process_dataset(val_dataset, tokenizer)
+        train_dataset_list.append(train_dataset)
+        val_dataset_list.append(val_dataset)
+    train_dataset = ConcatDataset(train_dataset_list)
+    val_dataset = ConcatDataset(val_dataset_list)
+
+    # dataloaders
+    train_dataloader = DataLoader(
+        train_dataset,
+        batch_size=args.train_batch_size,
+        num_workers=args.num_proc,
+        collate_fn=DefaultDataCollator(),
+        shuffle=True,
+    )
+    val_dataloader = DataLoader(
+        val_dataset,
+        batch_size=args.val_batch_size,
+        num_workers=args.num_proc,
+        collate_fn=DefaultDataCollator(),
+    )
+
+    # trainer
+    torch.set_float32_matmul_precision('medium')
+    if args.bf16:
+        precision = 'bf16-mixed'
+    elif args.fp16:
+        precision = '16-mixed'
+    else:
+        precision = "32-true"
+    lit_trainer = pl.Trainer(
+        accelerator='gpu',
+        precision=precision,
+        log_every_n_steps=5,
+        accumulate_grad_batches=32,
+        strategy=strategies.DeepSpeedStrategy(stage=3),
+    )
+    lit_trainer.fit(
+        lit_module,
+        train_dataloaders=train_dataloader,
+        val_dataloaders=val_dataloader,
+        ckpt_path=args.resume_from_ckpt_path,
+    )