aapot commited on
Commit
b221124
1 Parent(s): baf633c

Initial commit

Browse files
added_tokens.json ADDED
@@ -0,0 +1 @@
 
 
1
+ {"<|endoftext|>": 50257}
config.json ADDED
@@ -0,0 +1,38 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_name_or_path": "gpt2-large",
3
+ "activation_function": "gelu_new",
4
+ "architectures": [
5
+ "GPT2LMHeadModel"
6
+ ],
7
+ "attn_pdrop": 0.1,
8
+ "bos_token_id": 50256,
9
+ "embd_pdrop": 0.1,
10
+ "eos_token_id": 50256,
11
+ "initializer_range": 0.02,
12
+ "layer_norm_epsilon": 1e-05,
13
+ "model_type": "gpt2",
14
+ "n_ctx": 1024,
15
+ "n_embd": 1280,
16
+ "n_head": 20,
17
+ "n_inner": null,
18
+ "n_layer": 36,
19
+ "n_positions": 1024,
20
+ "reorder_and_upcast_attn": false,
21
+ "resid_pdrop": 0.1,
22
+ "scale_attn_by_inverse_layer_idx": false,
23
+ "scale_attn_weights": true,
24
+ "summary_activation": null,
25
+ "summary_first_dropout": 0.1,
26
+ "summary_proj_to_labels": true,
27
+ "summary_type": "cls_index",
28
+ "summary_use_proj": true,
29
+ "task_specific_params": {
30
+ "text-generation": {
31
+ "do_sample": true,
32
+ "max_length": 50
33
+ }
34
+ },
35
+ "transformers_version": "4.17.0.dev0",
36
+ "use_cache": true,
37
+ "vocab_size": 50257
38
+ }
flax_model_to_pytorch.py ADDED
@@ -0,0 +1,27 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import AutoModelForCausalLM, FlaxAutoModelForCausalLM, AutoTokenizer
2
+ import torch
3
+ import numpy as np
4
+ import jax
5
+ import jax.numpy as jnp
6
+
7
+ def to_f32(t):
8
+ return jax.tree_map(lambda x: x.astype(jnp.float32) if x.dtype == jnp.bfloat16 else x, t)
9
+
10
+ jax.config.update('jax_platform_name', 'cpu')
11
+ MODEL_PATH = "./"
12
+ model = FlaxAutoModelForCausalLM.from_pretrained(MODEL_PATH)
13
+ model.params = to_f32(model.params)
14
+ model.save_pretrained(MODEL_PATH)
15
+
16
+ pt_model = AutoModelForCausalLM.from_pretrained(
17
+ MODEL_PATH, from_flax=True).to('cpu')
18
+
19
+ input_ids = np.asarray(2 * [128 * [0]], dtype=np.int32)
20
+ input_ids_pt = torch.tensor(input_ids)
21
+
22
+ logits_pt = pt_model(input_ids_pt).logits
23
+ print(logits_pt)
24
+ logits_fx = model(input_ids).logits
25
+ print(logits_fx)
26
+
27
+ pt_model.save_pretrained(MODEL_PATH)
merges.txt ADDED
The diff for this file is too large to render. See raw diff
 
replace_token_script.py ADDED
@@ -0,0 +1,80 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ''''This script was used to replace the final index of tokenizer.json and vocab.json
2
+ with "<|endoftext|>" token. Also reassociate the corresponding merges'''
3
+
4
+ import json
5
+
6
+ tokenizer_path = 'tokenizer.json'
7
+ model_config_path = 'config.json'
8
+ vocab_path = 'vocab.json'
9
+
10
+ with open(vocab_path, "r") as f:
11
+ vocab_data = json.load(f)
12
+
13
+ with open(tokenizer_path, "r") as f:
14
+ tokenizer_data = json.load(f)
15
+
16
+ with open(model_config_path, "r") as f:
17
+ model_config = json.load(f)
18
+
19
+ model_vocab_size = model_config['vocab_size']
20
+ tokenizer_vocab = tokenizer_data['model']['vocab']
21
+
22
+ mergeslength = len(tokenizer_data['model']['merges'])
23
+
24
+ #readjust added_tokens 'id' to model_vocab_size - 1
25
+ tokenizer_data['added_tokens'][-1]['id'] = model_vocab_size - 1
26
+
27
+ final_index = model_vocab_size - 1
28
+ eos = '<|endoftext|>'
29
+
30
+ #retrieve the key of final index
31
+ old_key_final_index_tokenizer = list(tokenizer_data['model']['vocab'].keys())[final_index]
32
+ old_key_final_index_vocab = list(vocab_data.keys())[final_index]
33
+ old_key_final_index_vocab_min2 = list(vocab_data.keys())[final_index - 1]
34
+ old_key_final_index_tokenizer_merges = tokenizer_data['model']['merges'][mergeslength - 1]
35
+
36
+ print(f"old_key_final_index_tokenizer = {old_key_final_index_tokenizer}")
37
+ print(f"old_key_final_index_vocab = {old_key_final_index_vocab}")
38
+ print(f"old_key_final_index_vocab_min2 = {old_key_final_index_vocab_min2}")
39
+ print(f"old_key_final_index_tokenizer_merges = {old_key_final_index_tokenizer_merges}")
40
+
41
+ #replace old key with new key
42
+ tokenizer_data['model']['vocab']['<|endoftext|>'] = tokenizer_data['model']['vocab'][old_key_final_index_tokenizer]
43
+ vocab_data[eos] = vocab_data[old_key_final_index_vocab]
44
+
45
+ #replace the final merges idx with vocab_data - 1
46
+ tokenizer_data['model']['merges'] = tokenizer_data['model']['merges'][: mergeslength - 1]
47
+
48
+
49
+ #delete old key
50
+ del tokenizer_data['model']['vocab'][old_key_final_index_tokenizer]
51
+ del vocab_data[old_key_final_index_vocab]
52
+
53
+ #check updated key
54
+ old_key_final_index_tokenizer = list(tokenizer_data['model']['vocab'].keys())[final_index]
55
+ old_key_final_index_vocab = list(vocab_data.keys())[final_index]
56
+ old_key_final_index_tokenizer_merges = tokenizer_data['model']['merges'][mergeslength - 2]
57
+
58
+ print(len(tokenizer_data['model']['merges']))
59
+ print()
60
+ print(f"updated old_key_final_index_tokenizer = {old_key_final_index_tokenizer}")
61
+ print(f"updated old_key_final_index_vocab = {old_key_final_index_vocab}")
62
+ print(f"updated old_key_final_index_tokenizer_merges = {old_key_final_index_tokenizer_merges}")
63
+
64
+ with open(tokenizer_path, "w")as f:
65
+ json.dump(tokenizer_data, f)
66
+
67
+ with open(vocab_path, "w")as f:
68
+ json.dump(vocab_data, f)
69
+
70
+ with open('merges.txt') as f:
71
+ lines = f.readlines()
72
+
73
+ with open("merges.txt", "w") as f:
74
+ for i in range(len(lines) - 1):
75
+ f.write(lines[i])
76
+
77
+ with open('merges.txt') as f:
78
+ newlines = f.readlines()
79
+
80
+ print(f"newlines[len(newlines) - 1] = {newlines[len(newlines) - 1]}")
run_clm_flax.py ADDED
@@ -0,0 +1,892 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python
2
+ # coding=utf-8
3
+ # Copyright 2021 The HuggingFace Team All rights reserved.
4
+ #
5
+ # Licensed under the Apache License, Version 2.0 (the "License");
6
+ # you may not use this file except in compliance with the License.
7
+ # You may obtain a copy of the License at
8
+ #
9
+ # http://www.apache.org/licenses/LICENSE-2.0
10
+ #
11
+ # Unless required by applicable law or agreed to in writing, software
12
+ # distributed under the License is distributed on an "AS IS" BASIS,
13
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14
+ # See the License for the specific language governing permissions and
15
+ # limitations under the License.
16
+ """
17
+ Pre-training/Fine-tuning the library models for causal language modeling (GPT, GPT-2, CTRL, ...) on a text file or a dataset.
18
+
19
+ Here is the full list of checkpoints on the hub that can be fine-tuned by this script:
20
+ https://huggingface.co/models?filter=text-generation
21
+ """
22
+ # You can also adapt this script on your own causal language modeling task. Pointers for this are left as comments.
23
+
24
+ import json
25
+ import logging
26
+ import math
27
+ import os
28
+ import sys
29
+ import time
30
+ import gc
31
+ from dataclasses import asdict, dataclass, field
32
+ from enum import Enum
33
+ from itertools import chain
34
+ from pathlib import Path
35
+ from typing import Callable, Optional
36
+
37
+ import datasets
38
+ import numpy as np
39
+ from datasets import Dataset, load_dataset, load_from_disk
40
+ from tqdm import tqdm
41
+
42
+ import jax
43
+ import jax.numpy as jnp
44
+ import optax
45
+ import transformers
46
+ from flax import jax_utils, traverse_util
47
+ from flax.jax_utils import unreplicate
48
+ from flax.training import train_state
49
+ from flax.training.common_utils import get_metrics, onehot, shard, shard_prng_key
50
+ from huggingface_hub import Repository
51
+ from transformers import (
52
+ CONFIG_MAPPING,
53
+ FLAX_MODEL_FOR_CAUSAL_LM_MAPPING,
54
+ AutoConfig,
55
+ AutoTokenizer,
56
+ FlaxAutoModelForCausalLM,
57
+ HfArgumentParser,
58
+ is_tensorboard_available,
59
+ set_seed,
60
+ )
61
+ from transformers.file_utils import get_full_repo_name
62
+ from transformers.testing_utils import CaptureLogger
63
+
64
+ from distributed_shampoo import distributed_shampoo, GraftingType
65
+
66
+
67
+ logger = logging.getLogger(__name__)
68
+
69
+ MODEL_CONFIG_CLASSES = list(FLAX_MODEL_FOR_CAUSAL_LM_MAPPING.keys())
70
+ MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)
71
+
72
+
73
+ @dataclass
74
+ class TrainingArguments:
75
+ output_dir: str = field(
76
+ metadata={"help": "The output directory where the model predictions and checkpoints will be written."},
77
+ )
78
+ overwrite_output_dir: bool = field(
79
+ default=False,
80
+ metadata={
81
+ "help": (
82
+ "Overwrite the content of the output directory. "
83
+ "Use this to continue training if output_dir points to a checkpoint directory."
84
+ )
85
+ },
86
+ )
87
+ do_train: bool = field(default=False, metadata={"help": "Whether to run training."})
88
+ do_eval: bool = field(default=False, metadata={"help": "Whether to run eval on the dev set."})
89
+ per_device_train_batch_size: int = field(
90
+ default=8, metadata={"help": "Batch size per GPU/TPU core/CPU for training."}
91
+ )
92
+ per_device_eval_batch_size: int = field(
93
+ default=8, metadata={"help": "Batch size per GPU/TPU core/CPU for evaluation."}
94
+ )
95
+ gradient_accumulation_steps: int = field(
96
+ default=1,
97
+ metadata={
98
+ "help": "Number of updates steps to accumulate before performing a backward/update pass."
99
+ },
100
+ )
101
+ learning_rate: float = field(default=5e-5, metadata={"help": "The initial learning rate for AdamW."})
102
+ weight_decay: float = field(default=0.0, metadata={"help": "Weight decay for AdamW if we apply some."})
103
+ adam_beta1: float = field(default=0.9, metadata={"help": "Beta1 for AdamW optimizer"})
104
+ adam_beta2: float = field(default=0.999, metadata={"help": "Beta2 for AdamW optimizer"})
105
+ adam_epsilon: float = field(default=1e-8, metadata={"help": "Epsilon for AdamW optimizer."})
106
+ adafactor: bool = field(default=False, metadata={"help": "Whether or not to replace AdamW by Adafactor."})
107
+ distributed_shampoo: bool = field(
108
+ default=False, metadata={"help": "Use Distributed Shampoo optimizer instead of AdamW."},
109
+ )
110
+ quantize_shampoo: bool = field(
111
+ default=False, metadata={"help": "Quantize Distributed Shampoo optimizer."},
112
+ )
113
+ num_train_epochs: float = field(default=3.0, metadata={"help": "Total number of training epochs to perform."})
114
+ warmup_steps: int = field(default=0, metadata={"help": "Linear warmup over warmup_steps."})
115
+ warmup_ratio: float = field(default=0.0, metadata={"help": "Linear warmup ratio of total train steps."})
116
+ cosine_decay: bool = field(
117
+ default=False, metadata={"help": "Whether or not to use cosine decay instead of the basic linear decay schedule."}
118
+ )
119
+ gradient_clipping: bool = field(
120
+ default=False, metadata={"help": "Whether or not to use gradient clipping."}
121
+ )
122
+ logging_steps: int = field(default=500, metadata={"help": "Log every X updates steps."})
123
+ save_steps: int = field(default=500, metadata={"help": "Save checkpoint every X updates steps."})
124
+ eval_steps: int = field(default=None, metadata={"help": "Run an evaluation every X steps."})
125
+ seed: int = field(default=42, metadata={"help": "Random seed that will be set at the beginning of training."})
126
+ push_to_hub: bool = field(
127
+ default=False, metadata={"help": "Whether or not to upload the trained model to the model hub after training."}
128
+ )
129
+ hub_model_id: str = field(
130
+ default=None, metadata={"help": "The name of the repository to keep in sync with the local `output_dir`."}
131
+ )
132
+ hub_token: str = field(default=None, metadata={"help": "The token to use to push to the Model Hub."})
133
+
134
+ def __post_init__(self):
135
+ if self.output_dir is not None:
136
+ self.output_dir = os.path.expanduser(self.output_dir)
137
+
138
+ def to_dict(self):
139
+ """
140
+ Serializes this instance while replace `Enum` by their values (for JSON serialization support). It obfuscates
141
+ the token values by removing their value.
142
+ """
143
+ d = asdict(self)
144
+ for k, v in d.items():
145
+ if isinstance(v, Enum):
146
+ d[k] = v.value
147
+ if isinstance(v, list) and len(v) > 0 and isinstance(v[0], Enum):
148
+ d[k] = [x.value for x in v]
149
+ if k.endswith("_token"):
150
+ d[k] = f"<{k.upper()}>"
151
+ return d
152
+
153
+
154
+ @dataclass
155
+ class ModelArguments:
156
+ """
157
+ Arguments pertaining to which model/config/tokenizer we are going to fine-tune, or train from scratch.
158
+ """
159
+
160
+ model_name_or_path: Optional[str] = field(
161
+ default=None,
162
+ metadata={
163
+ "help": "The model checkpoint for weights initialization."
164
+ "Don't set if you want to train a model from scratch."
165
+ },
166
+ )
167
+ model_type: Optional[str] = field(
168
+ default=None,
169
+ metadata={"help": "If training from scratch, pass a model type from the list: " + ", ".join(MODEL_TYPES)},
170
+ )
171
+ config_name: Optional[str] = field(
172
+ default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
173
+ )
174
+ tokenizer_name: Optional[str] = field(
175
+ default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
176
+ )
177
+ cache_dir: Optional[str] = field(
178
+ default=None, metadata={"help": "Where do you want to store the pretrained models downloaded from s3"}
179
+ )
180
+ use_fast_tokenizer: bool = field(
181
+ default=True,
182
+ metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."},
183
+ )
184
+ dtype: Optional[str] = field(
185
+ default="float32",
186
+ metadata={
187
+ "help": "Floating-point format in which the model weights should be initialized and trained. Choose one of `[float32, float16, bfloat16]`."
188
+ },
189
+ )
190
+
191
+
192
+ @dataclass
193
+ class DataTrainingArguments:
194
+ """
195
+ Arguments pertaining to what data we are going to input our model for training and eval.
196
+ """
197
+
198
+ dataset_name: Optional[str] = field(
199
+ default=None, metadata={"help": "The name of the dataset to use (via the datasets library)."}
200
+ )
201
+ dataset_config_name: Optional[str] = field(
202
+ default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
203
+ )
204
+ dataset_filepath: Optional[str] = field(
205
+ default=None, metadata={"help": "Filepath to locally saved HF Dataset (with 'dataset.save_to_disk' method) to use for training"}
206
+ )
207
+ train_file: Optional[str] = field(default=None, metadata={"help": "The input training data file (a text file)."})
208
+ validation_file: Optional[str] = field(
209
+ default=None,
210
+ metadata={"help": "An optional input evaluation data file to evaluate the perplexity on (a text file)."},
211
+ )
212
+ max_train_samples: Optional[int] = field(
213
+ default=None,
214
+ metadata={
215
+ "help": "For debugging purposes or quicker training, truncate the number of training examples to this "
216
+ "value if set."
217
+ },
218
+ )
219
+ max_eval_samples: Optional[int] = field(
220
+ default=None,
221
+ metadata={
222
+ "help": "For debugging purposes or quicker training, truncate the number of evaluation examples to this "
223
+ "value if set."
224
+ },
225
+ )
226
+ overwrite_cache: bool = field(
227
+ default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
228
+ )
229
+ validation_split_percentage: Optional[int] = field(
230
+ default=5,
231
+ metadata={
232
+ "help": "The percentage of the train set used as validation set in case there's no validation split"
233
+ },
234
+ )
235
+ block_size: Optional[int] = field(
236
+ default=None,
237
+ metadata={
238
+ "help": "Optional input sequence length after tokenization. "
239
+ "The training dataset will be truncated in block of this size for training. "
240
+ "Default to the model max input length for single sentence inputs (take into account special tokens)."
241
+ },
242
+ )
243
+ overwrite_cache: bool = field(
244
+ default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
245
+ )
246
+ preprocessing_num_workers: Optional[int] = field(
247
+ default=None,
248
+ metadata={"help": "The number of processes to use for the preprocessing."},
249
+ )
250
+ keep_linebreaks: bool = field(
251
+ default=True, metadata={"help": "Whether to keep line breaks when using TXT files or not."}
252
+ )
253
+
254
+ def __post_init__(self):
255
+ if self.dataset_name is None and self.train_file is None and self.dataset_filepath is None and self.validation_file is None:
256
+ raise ValueError("Need either a dataset name or a training/validation file.")
257
+ else:
258
+ if self.train_file is not None:
259
+ extension = self.train_file.split(".")[-1]
260
+ assert extension in ["csv", "json", "txt"], "`train_file` should be a csv, a json or a txt file."
261
+ if self.validation_file is not None:
262
+ extension = self.validation_file.split(".")[-1]
263
+ assert extension in ["csv", "json", "txt"], "`validation_file` should be a csv, a json or a txt file."
264
+
265
+
266
+ class TrainState(train_state.TrainState):
267
+ dropout_rng: jnp.ndarray
268
+
269
+ def replicate(self):
270
+ return jax_utils.replicate(self).replace(dropout_rng=shard_prng_key(self.dropout_rng))
271
+
272
+
273
+ def data_loader(rng: jax.random.PRNGKey, dataset: Dataset, batch_size: int, shuffle: bool = False):
274
+ """
275
+ Returns batches of size `batch_size` from truncated `dataset`, sharded over all local devices.
276
+ Shuffle batches if `shuffle` is `True`.
277
+ """
278
+ steps_per_epoch = len(dataset) // batch_size
279
+
280
+ if shuffle:
281
+ batch_idx = jax.random.permutation(rng, len(dataset))
282
+ else:
283
+ batch_idx = jnp.arange(len(dataset))
284
+
285
+ batch_idx = batch_idx[: steps_per_epoch * batch_size] # Skip incomplete batch.
286
+ batch_idx = batch_idx.reshape((steps_per_epoch, batch_size))
287
+
288
+ for idx in batch_idx:
289
+ batch = dataset[idx]
290
+ batch = {k: np.array(v) for k, v in batch.items()}
291
+
292
+ yield batch
293
+
294
+
295
+ def write_train_metric(summary_writer, train_metrics, train_time, step):
296
+ summary_writer.scalar("train_time", train_time, step)
297
+
298
+ train_metrics = get_metrics(train_metrics)
299
+ for key, vals in train_metrics.items():
300
+ tag = f"train_{key}"
301
+ for i, val in enumerate(vals):
302
+ summary_writer.scalar(tag, val, step - len(vals) + i + 1)
303
+
304
+
305
+ def write_eval_metric(summary_writer, eval_metrics, step):
306
+ for metric_name, value in eval_metrics.items():
307
+ summary_writer.scalar(f"eval_{metric_name}", value, step)
308
+
309
+
310
+ def create_linear_learning_rate_fn(
311
+ train_ds_size: int, train_batch_size: int, num_train_epochs: int, num_warmup_steps: int, learning_rate: float
312
+ ) -> Callable[[int], jnp.array]:
313
+ """Returns a linear warmup, linear decay learning rate function."""
314
+ steps_per_epoch = train_ds_size // train_batch_size
315
+ num_train_steps = steps_per_epoch * num_train_epochs
316
+ warmup_fn = optax.linear_schedule(init_value=0.0, end_value=learning_rate, transition_steps=num_warmup_steps)
317
+ decay_fn = optax.linear_schedule(
318
+ init_value=learning_rate, end_value=0, transition_steps=num_train_steps - num_warmup_steps
319
+ )
320
+ schedule_fn = optax.join_schedules(schedules=[warmup_fn, decay_fn], boundaries=[num_warmup_steps])
321
+ return schedule_fn
322
+
323
+ def create_cosine_learning_rate_fn(
324
+ train_ds_size: int, train_batch_size: int, num_train_epochs: int, num_warmup_steps: int, learning_rate: float
325
+ ) -> Callable[[int], jnp.array]:
326
+ """Returns a linear warmup, cosine decay learning rate function."""
327
+ steps_per_epoch = train_ds_size // train_batch_size
328
+ num_train_steps = steps_per_epoch * num_train_epochs
329
+ warmup_fn = optax.linear_schedule(init_value=0.0, end_value=learning_rate, transition_steps=num_warmup_steps)
330
+ decay_fn = optax.cosine_decay_schedule(
331
+ init_value=learning_rate, decay_steps=num_train_steps - num_warmup_steps, alpha=0.1
332
+ )
333
+ schedule_fn = optax.join_schedules(schedules=[warmup_fn, decay_fn], boundaries=[num_warmup_steps])
334
+ return schedule_fn
335
+
336
+
337
+ def main():
338
+ # See all possible arguments in src/transformers/training_args.py
339
+ # or by passing the --help flag to this script.
340
+ # We now keep distinct sets of args, for a cleaner separation of concerns.
341
+
342
+ parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
343
+ if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
344
+ # If we pass only one argument to the script and it's the path to a json file,
345
+ # let's parse it to get our arguments.
346
+ model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
347
+ else:
348
+ model_args, data_args, training_args = parser.parse_args_into_dataclasses()
349
+
350
+ if (
351
+ os.path.exists(training_args.output_dir)
352
+ and os.listdir(training_args.output_dir)
353
+ and training_args.do_train
354
+ and not training_args.overwrite_output_dir
355
+ ):
356
+ raise ValueError(
357
+ f"Output directory ({training_args.output_dir}) already exists and is not empty."
358
+ "Use --overwrite_output_dir to overcome."
359
+ )
360
+
361
+ # Make one log on every process with the configuration for debugging.
362
+ logging.basicConfig(
363
+ format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
364
+ datefmt="%m/%d/%Y %H:%M:%S",
365
+ level=logging.INFO,
366
+ )
367
+ # Setup logging, we only want one process per machine to log things on the screen.
368
+ logger.setLevel(logging.INFO if jax.process_index() == 0 else logging.ERROR)
369
+ if jax.process_index() == 0:
370
+ datasets.utils.logging.set_verbosity_warning()
371
+ transformers.utils.logging.set_verbosity_info()
372
+ else:
373
+ datasets.utils.logging.set_verbosity_error()
374
+ transformers.utils.logging.set_verbosity_error()
375
+
376
+ # Set the verbosity to info of the Transformers logger (on main process only):
377
+ logger.info(f"Training/evaluation parameters {training_args}")
378
+
379
+ # Set seed before initializing model.
380
+ set_seed(training_args.seed)
381
+
382
+ # Handle the repository creation
383
+ if training_args.push_to_hub:
384
+ if training_args.hub_model_id is None:
385
+ repo_name = get_full_repo_name(
386
+ Path(training_args.output_dir).absolute().name, token=training_args.hub_token
387
+ )
388
+ else:
389
+ repo_name = training_args.hub_model_id
390
+ repo = Repository(training_args.output_dir, clone_from=repo_name)
391
+
392
+ # Get the datasets: you can either provide your own CSV/JSON/TXT training and evaluation files (see below)
393
+ # or just provide the name of one of the public datasets available on the hub at https://huggingface.co/datasets/
394
+ # (the dataset will be downloaded automatically from the datasets Hub).
395
+ #
396
+ # For CSV/JSON files, this script will use the column called 'text' or the first column if no column called
397
+ # 'text' is found. You can easily tweak this behavior (see below).
398
+ #
399
+ # In distributed training, the load_dataset function guarantees that only one local process can concurrently
400
+ # download the dataset.
401
+ if data_args.dataset_name is not None:
402
+ # Downloading and loading a dataset from the hub.
403
+ dataset = load_dataset(
404
+ data_args.dataset_name, data_args.dataset_config_name, cache_dir=model_args.cache_dir, keep_in_memory=False
405
+ )
406
+
407
+ if "validation" not in dataset.keys():
408
+ dataset["validation"] = load_dataset(
409
+ data_args.dataset_name,
410
+ data_args.dataset_config_name,
411
+ split=f"train[:{data_args.validation_split_percentage}%]",
412
+ cache_dir=model_args.cache_dir,
413
+ )
414
+ dataset["train"] = load_dataset(
415
+ data_args.dataset_name,
416
+ data_args.dataset_config_name,
417
+ split=f"train[{data_args.validation_split_percentage}%:]",
418
+ cache_dir=model_args.cache_dir,
419
+ )
420
+
421
+ elif data_args.dataset_filepath is not None:
422
+ # Loading a dataset from local file.
423
+ dataset = load_from_disk(data_args.dataset_filepath)
424
+ if "validation" not in dataset.keys():
425
+ dataset = datasets.train_test_split(test_size=data_args.validation_split_percentage/100)
426
+ dataset["validation"] = dataset["test"]
427
+ del dataset["test"]
428
+
429
+ else:
430
+ data_files = {}
431
+ dataset_args = {}
432
+ if data_args.train_file is not None:
433
+ data_files["train"] = data_args.train_file
434
+ if data_args.validation_file is not None:
435
+ data_files["validation"] = data_args.validation_file
436
+ extension = data_args.train_file.split(".")[-1]
437
+ if extension == "txt":
438
+ extension = "text"
439
+ dataset_args["keep_linebreaks"] = data_args.keep_linebreaks
440
+ dataset = load_dataset(extension, data_files=data_files, cache_dir=model_args.cache_dir, **dataset_args)
441
+
442
+ if "validation" not in dataset.keys():
443
+ dataset["validation"] = load_dataset(
444
+ extension,
445
+ data_files=data_files,
446
+ split=f"train[:{data_args.validation_split_percentage}%]",
447
+ cache_dir=model_args.cache_dir,
448
+ **dataset_args,
449
+ )
450
+ dataset["train"] = load_dataset(
451
+ extension,
452
+ data_files=data_files,
453
+ split=f"train[{data_args.validation_split_percentage}%:]",
454
+ cache_dir=model_args.cache_dir,
455
+ **dataset_args,
456
+ )
457
+ # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at
458
+ # https://huggingface.co/docs/datasets/loading_datasets.html.
459
+
460
+ # Load pretrained model and tokenizer
461
+
462
+ # Distributed training:
463
+ # The .from_pretrained methods guarantee that only one local process can concurrently
464
+ # download model & vocab.
465
+ if model_args.config_name:
466
+ config = AutoConfig.from_pretrained(model_args.config_name, cache_dir=model_args.cache_dir)
467
+ elif model_args.model_name_or_path:
468
+ config = AutoConfig.from_pretrained(model_args.model_name_or_path, cache_dir=model_args.cache_dir)
469
+ else:
470
+ config = CONFIG_MAPPING[model_args.model_type]()
471
+ logger.warning("You are instantiating a new config instance from scratch.")
472
+
473
+ if model_args.tokenizer_name:
474
+ tokenizer = AutoTokenizer.from_pretrained(
475
+ model_args.tokenizer_name, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_tokenizer
476
+ )
477
+ elif model_args.model_name_or_path:
478
+ tokenizer = AutoTokenizer.from_pretrained(
479
+ model_args.model_name_or_path, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_tokenizer
480
+ )
481
+ else:
482
+ raise ValueError(
483
+ "You are instantiating a new tokenizer from scratch. This is not supported by this script."
484
+ "You can do it from another script, save it, and load it from here, using --tokenizer_name."
485
+ )
486
+
487
+ if tokenizer.pad_token is None:
488
+ tokenizer.pad_token = tokenizer.eos_token
489
+
490
+ if model_args.model_name_or_path:
491
+ model = FlaxAutoModelForCausalLM.from_pretrained(
492
+ model_args.model_name_or_path, config=config, seed=training_args.seed, dtype=getattr(jnp, model_args.dtype)
493
+ )
494
+ else:
495
+ model = FlaxAutoModelForCausalLM.from_config(
496
+ config, seed=training_args.seed, dtype=getattr(jnp, model_args.dtype)
497
+ )
498
+
499
+ # Preprocessing the datasets.
500
+ # First we tokenize all the texts.
501
+ if training_args.do_train:
502
+ column_names = dataset["train"].column_names
503
+ else:
504
+ column_names = dataset["validation"].column_names
505
+ text_column_name = "text" if "text" in column_names else column_names[0]
506
+
507
+ # since this will be pickled to avoid _LazyModule error in Hasher force logger loading before tokenize_function
508
+ tok_logger = transformers.utils.logging.get_logger("transformers.tokenization_utils_base")
509
+
510
+ def tokenize_function(examples):
511
+ with CaptureLogger(tok_logger) as cl:
512
+ output = tokenizer(examples[text_column_name])
513
+ # clm input could be much much longer than block_size
514
+ if "Token indices sequence length is longer than the" in cl.out:
515
+ tok_logger.warning(
516
+ "^^^^^^^^^^^^^^^^ Please ignore the warning above - this long input will be chunked into smaller bits before being passed to the model."
517
+ )
518
+ return output
519
+
520
+ tokenized_datasets = dataset.map(
521
+ tokenize_function,
522
+ batched=True,
523
+ num_proc=data_args.preprocessing_num_workers,
524
+ remove_columns=column_names,
525
+ load_from_cache_file=not data_args.overwrite_cache,
526
+ )
527
+
528
+ if data_args.block_size is None:
529
+ block_size = tokenizer.model_max_length
530
+ if block_size > config.max_position_embeddings:
531
+ logger.warning(
532
+ f"The tokenizer picked seems to have a very large `model_max_length` ({tokenizer.model_max_length}). "
533
+ "Picking 1024 instead. You can change that default value by passing --block_size xxx."
534
+ )
535
+ block_size = 1024
536
+ else:
537
+ if data_args.block_size > tokenizer.model_max_length:
538
+ logger.warning(
539
+ f"The block_size passed ({data_args.block_size}) is larger than the maximum length for the model"
540
+ f"({tokenizer.model_max_length}). Using block_size={tokenizer.model_max_length}."
541
+ )
542
+ block_size = min(data_args.block_size, tokenizer.model_max_length)
543
+
544
+ # Main data processing function that will concatenate all texts from our dataset and generate chunks of block_size.
545
+ def group_texts(examples):
546
+ # Concatenate all texts.
547
+ concatenated_examples = {k: list(chain(*examples[k])) for k in examples.keys()}
548
+ total_length = len(concatenated_examples[list(examples.keys())[0]])
549
+ # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can
550
+ # customize this part to your needs.
551
+ if total_length >= block_size:
552
+ total_length = (total_length // block_size) * block_size
553
+ # Split by chunks of max_len.
554
+ result = {
555
+ k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
556
+ for k, t in concatenated_examples.items()
557
+ }
558
+ result["labels"] = result["input_ids"].copy()
559
+ return result
560
+
561
+ # Note that with `batched=True`, this map processes 1,000 texts together, so group_texts throws away a remainder
562
+ # for each of those groups of 1,000 texts. You can adjust that batch_size here but a higher value might be slower
563
+ # to preprocess.
564
+ #
565
+ # To speed up this part, we use multiprocessing. See the documentation of the map method for more information:
566
+ # https://huggingface.co/docs/datasets/package_reference/main_classes.html#datasets.Dataset.map
567
+
568
+ lm_datasets = tokenized_datasets.map(
569
+ group_texts,
570
+ batched=True,
571
+ num_proc=data_args.preprocessing_num_workers,
572
+ load_from_cache_file=not data_args.overwrite_cache,
573
+ )
574
+
575
+ if training_args.do_train:
576
+ if "train" not in tokenized_datasets:
577
+ raise ValueError("--do_train requires a train dataset")
578
+ train_dataset = lm_datasets["train"]
579
+ if data_args.max_train_samples is not None:
580
+ train_dataset = train_dataset.select(range(data_args.max_train_samples))
581
+
582
+ # test to see that tokenization worked
583
+ detokenized_example = tokenizer.decode(train_dataset[0]["input_ids"])
584
+ logger.info(f"Detokenized example: {detokenized_example}")
585
+ detokenized_example = tokenizer.decode(train_dataset[-1]["input_ids"])
586
+ logger.info(f"Detokenized example 2: {detokenized_example}")
587
+
588
+ if training_args.do_eval:
589
+ if "validation" not in tokenized_datasets:
590
+ raise ValueError("--do_eval requires a validation dataset")
591
+ eval_dataset = lm_datasets["validation"]
592
+ if data_args.max_eval_samples is not None:
593
+ eval_dataset = eval_dataset.select(range(data_args.max_eval_samples))
594
+
595
+ # Enable tensorboard only on the master node
596
+ has_tensorboard = is_tensorboard_available()
597
+ if has_tensorboard and jax.process_index() == 0:
598
+ try:
599
+ from flax.metrics.tensorboard import SummaryWriter
600
+
601
+ summary_writer = SummaryWriter(log_dir=Path(os.path.join(training_args.output_dir, "runs")))
602
+ except ImportError as ie:
603
+ has_tensorboard = False
604
+ logger.warning(
605
+ f"Unable to display metrics through TensorBoard because some package are not installed: {ie}"
606
+ )
607
+ else:
608
+ logger.warning(
609
+ "Unable to display metrics through TensorBoard because the package is not installed: "
610
+ "Please run pip install tensorboard to enable."
611
+ )
612
+
613
+ # Initialize our training
614
+ rng = jax.random.PRNGKey(training_args.seed)
615
+ rng, dropout_rng = jax.random.split(rng)
616
+
617
+ # Store some constant
618
+ num_epochs = int(training_args.num_train_epochs)
619
+ train_batch_size = int(training_args.per_device_train_batch_size) * jax.device_count() * training_args.gradient_accumulation_steps
620
+ eval_batch_size = int(training_args.per_device_eval_batch_size) * jax.device_count()
621
+ steps_per_epoch = len(train_dataset) // train_batch_size
622
+ total_train_steps = steps_per_epoch * num_epochs
623
+
624
+ if training_args.warmup_ratio > 0:
625
+ warmup_steps = int(total_train_steps * training_args.warmup_ratio)
626
+ else:
627
+ warmup_steps = training_args.warmup_steps
628
+
629
+ # Create learning rate schedule
630
+ if training_args.cosine_decay:
631
+ lr_schedule_fn = create_cosine_learning_rate_fn(
632
+ len(train_dataset),
633
+ train_batch_size,
634
+ training_args.num_train_epochs,
635
+ warmup_steps,
636
+ training_args.learning_rate,
637
+ )
638
+ else:
639
+ lr_schedule_fn = create_linear_learning_rate_fn(
640
+ len(train_dataset),
641
+ train_batch_size,
642
+ training_args.num_train_epochs,
643
+ warmup_steps,
644
+ training_args.learning_rate,
645
+ )
646
+
647
+ # We use Optax's "masking" functionality to not apply weight decay
648
+ # to bias and LayerNorm scale parameters. decay_mask_fn returns a
649
+ # mask boolean with the same structure as the parameters.
650
+ # The mask is True for parameters that should be decayed.
651
+ # Note that this mask is specifically adapted for FlaxGPT2.
652
+ # For other models, one should correct the layer norm parameter naming
653
+ # accordingly.
654
+ def decay_mask_fn(params):
655
+ flat_params = traverse_util.flatten_dict(params)
656
+ flat_mask = {
657
+ path: (path[-1] != "bias" and path[-2:] not in [("ln_1", "scale"), ("ln_2", "scale"), ("ln_f", "scale")])
658
+ for path in flat_params
659
+ }
660
+ return traverse_util.unflatten_dict(flat_mask)
661
+
662
+ # create adam optimizer
663
+ if training_args.adafactor:
664
+ # We use the default parameters here to initialize adafactor,
665
+ # For more details about the parameters please check https://github.com/deepmind/optax/blob/ed02befef9bf81cbbf236be3d2b0e032e9ed4a40/optax/_src/alias.py#L74
666
+ optimizer = optax.adafactor(
667
+ learning_rate=lr_schedule_fn,
668
+ )
669
+
670
+ elif training_args.distributed_shampoo:
671
+ # parameters from https://github.com/tensorflow/lingvo/blob/03ee9d7cd50764b0424c7c863733c91fc0b053ec/lingvo/jax/optimizers.py#L729
672
+ # Notes:
673
+ # - mask for weight decay is not implemented but we don't use it anyway
674
+ optimizer = distributed_shampoo(
675
+ lr_schedule_fn,
676
+ block_size=1536, # recommended default for large LM is 1536
677
+ beta1=0.9,
678
+ beta2=0.999,
679
+ diagonal_epsilon=1e-10,
680
+ matrix_epsilon=1e-8,
681
+ weight_decay=0.0,
682
+ start_preconditioning_step=1001,
683
+ preconditioning_compute_steps=10,
684
+ statistics_compute_steps=1,
685
+ best_effort_shape_interpretation=True,
686
+ graft_type=GraftingType.RMSPROP_NORMALIZED,
687
+ nesterov=False,
688
+ exponent_override=0,
689
+ batch_axis_name="batch",
690
+ inverse_failure_threshold=0.1,
691
+ moving_average_for_momentum=True,
692
+ skip_preconditioning_dim_size_gt=4096,
693
+ clip_by_scaled_gradient_norm=None,
694
+ precision=jax.lax.Precision.HIGHEST,
695
+ best_effort_memory_usage_reduction=training_args.quantize_shampoo,
696
+ )
697
+ else:
698
+ optimizer = optax.adamw(
699
+ learning_rate=lr_schedule_fn,
700
+ b1=training_args.adam_beta1,
701
+ b2=training_args.adam_beta2,
702
+ eps=training_args.adam_epsilon,
703
+ weight_decay=training_args.weight_decay,
704
+ mask=decay_mask_fn,
705
+ )
706
+ if training_args.gradient_clipping:
707
+ optimizer = optax.chain(
708
+ optax.clip_by_global_norm(1.),
709
+ optimizer
710
+ )
711
+
712
+ # add gradient accumulation
713
+ if training_args.gradient_accumulation_steps > 1:
714
+ optimizer = optax.chain(
715
+ optax.apply_every(training_args.gradient_accumulation_steps), optimizer
716
+ )
717
+
718
+ # Setup train state
719
+ state = TrainState.create(apply_fn=model.__call__, params=model.params, tx=optimizer, dropout_rng=dropout_rng)
720
+
721
+ def loss_fn(logits, labels):
722
+ shift_logits = logits[..., :-1, :]
723
+ shift_labels = labels[..., 1:]
724
+ loss = optax.softmax_cross_entropy(shift_logits, onehot(shift_labels, shift_logits.shape[-1]))
725
+ return loss.mean()
726
+
727
+ # Define gradient update step fn
728
+ def train_step(state, batch):
729
+ dropout_rng, new_dropout_rng = jax.random.split(state.dropout_rng)
730
+
731
+ def compute_loss(params):
732
+ labels = batch.pop("labels")
733
+ logits = state.apply_fn(**batch, params=params, dropout_rng=dropout_rng, train=True)[0]
734
+ loss = loss_fn(logits, labels)
735
+ return loss
736
+
737
+ grad_fn = jax.value_and_grad(compute_loss)
738
+ loss, grad = grad_fn(state.params)
739
+ grad = jax.lax.pmean(grad, "batch")
740
+
741
+ new_state = state.apply_gradients(grads=grad, dropout_rng=new_dropout_rng)
742
+
743
+ metrics = {"loss": loss, "learning_rate": lr_schedule_fn(state.step // training_args.gradient_accumulation_steps)}
744
+ metrics = jax.lax.pmean(metrics, axis_name="batch")
745
+
746
+ return new_state, metrics
747
+
748
+ # Define eval fn
749
+ def eval_step(params, batch):
750
+ labels = batch.pop("labels")
751
+ logits = model(**batch, params=params, train=False)[0]
752
+ loss = loss_fn(logits, labels)
753
+
754
+ # summarize metrics
755
+ metrics = {"loss": loss}
756
+ metrics = jax.lax.pmean(metrics, axis_name="batch")
757
+ return metrics
758
+
759
+ # Create parallel version of the train and eval step
760
+ p_train_step = jax.pmap(train_step, "batch", donate_argnums=(0,))
761
+ p_eval_step = jax.pmap(eval_step, "batch")
762
+
763
+ # Replicate the train state on each device
764
+ state = state.replicate()
765
+
766
+ logger.info("***** Running training *****")
767
+ logger.info(f" Num examples = {len(train_dataset)}")
768
+ logger.info(f" Num Epochs = {num_epochs}")
769
+ logger.info(f" Instantaneous batch size per device = {training_args.per_device_train_batch_size}")
770
+ logger.info(f" Total train batch size (w. parallel & distributed) = {train_batch_size}")
771
+ logger.info(f" Total optimization steps = {total_train_steps}")
772
+
773
+ train_time = 0
774
+ train_metrics = []
775
+ epochs = tqdm(range(num_epochs), desc="Epoch ... ", position=0)
776
+ for epoch in epochs:
777
+ # ======================== Training ================================
778
+ train_start = time.time()
779
+
780
+ # Create sampling rng
781
+ rng, input_rng = jax.random.split(rng)
782
+
783
+ # Generate an epoch by shuffling sampling indices from the train dataset
784
+ train_loader = data_loader(input_rng, train_dataset, train_batch_size // training_args.gradient_accumulation_steps, shuffle=True)
785
+ steps_per_epoch = len(train_dataset) // train_batch_size
786
+ # train
787
+ steps_trained_progress_bar = tqdm(range(steps_per_epoch), desc="Training...", position=1,
788
+ leave=False)
789
+ for step in range(steps_per_epoch * training_args.gradient_accumulation_steps):
790
+ batch = next(train_loader)
791
+ batch = shard(batch)
792
+ state, train_metric = p_train_step(state, batch)
793
+ train_metrics.append(train_metric)
794
+
795
+ cur_step = epoch * (steps_per_epoch*training_args.gradient_accumulation_steps) + step
796
+
797
+ if step % training_args.gradient_accumulation_steps == 0:
798
+ steps_trained_progress_bar.update(1)
799
+
800
+ if cur_step % (training_args.logging_steps * training_args.gradient_accumulation_steps) == 0 and cur_step > 0:
801
+ # Save metrics
802
+ train_metric = unreplicate(train_metric)
803
+ train_time += time.time() - train_start
804
+ if has_tensorboard and jax.process_index() == 0:
805
+ write_train_metric(summary_writer, train_metrics, train_time, cur_step)
806
+
807
+ epochs.write(
808
+ f"Step... ({cur_step} | Loss: {train_metric['loss'].mean()}, Learning Rate: {train_metric['learning_rate'].mean()})"
809
+ )
810
+
811
+ train_metrics = []
812
+
813
+ if cur_step % (training_args.eval_steps * training_args.gradient_accumulation_steps) == 0 and cur_step > 0:
814
+ # ======================== Evaluating ==============================
815
+ eval_metrics = []
816
+ eval_loader = data_loader(input_rng, eval_dataset, eval_batch_size)
817
+ eval_steps = len(eval_dataset) // eval_batch_size
818
+ for _ in tqdm(range(eval_steps), desc="Evaluating...", position=2, leave=False):
819
+ # Model forward
820
+ batch = next(eval_loader)
821
+ batch = shard(batch)
822
+ metrics = p_eval_step(state.params, batch)
823
+ eval_metrics.append(metrics)
824
+
825
+ # normalize eval metrics
826
+ eval_metrics = get_metrics(eval_metrics)
827
+ eval_metrics = jax.tree_map(jnp.mean, eval_metrics)
828
+
829
+ try:
830
+ eval_metrics["perplexity"] = math.exp(eval_metrics["loss"])
831
+ except OverflowError:
832
+ eval_metrics["perplexity"] = float("inf")
833
+
834
+ # Print metrics and update progress bar
835
+ desc = f"Step... ({cur_step} | Eval Loss: {eval_metrics['loss']} | Eval Perplexity: {eval_metrics['perplexity']})"
836
+ epochs.write(desc)
837
+ epochs.desc = desc
838
+
839
+ # Save metrics
840
+ if has_tensorboard and jax.process_index() == 0:
841
+ write_eval_metric(summary_writer, eval_metrics, cur_step)
842
+
843
+ if cur_step % (training_args.save_steps * training_args.gradient_accumulation_steps) == 0 and cur_step > 0:
844
+ # save checkpoint after each epoch and push checkpoint to the hub
845
+ if jax.process_index() == 0:
846
+ params = jax.device_get(unreplicate(state.params))
847
+ model.save_pretrained(training_args.output_dir, params=params)
848
+ tokenizer.save_pretrained(training_args.output_dir)
849
+ if training_args.push_to_hub:
850
+ repo.push_to_hub(commit_message=f"Saving weights and logs of step {cur_step}", blocking=False, auto_lfs_prune=True)
851
+
852
+ # save also at the end of epoch
853
+ try:
854
+ if jax.process_index() == 0:
855
+ params = jax.device_get(unreplicate(state.params))
856
+ model.save_pretrained(training_args.output_dir, params=params)
857
+ tokenizer.save_pretrained(training_args.output_dir)
858
+ if training_args.push_to_hub:
859
+ repo.push_to_hub(commit_message=f"Saving weights and logs of epoch {epoch}", blocking=False, auto_lfs_prune=True)
860
+ except:
861
+ # push to hub fails the whole script if nothing new to commit
862
+ pass
863
+
864
+ # Eval after training
865
+ if training_args.do_eval:
866
+ eval_metrics = []
867
+ eval_loader = data_loader(input_rng, eval_dataset, eval_batch_size)
868
+ eval_steps = len(eval_dataset) // eval_batch_size
869
+ for _ in tqdm(range(eval_steps), desc="Evaluating...", position=2, leave=False):
870
+ # Model forward
871
+ batch = shard(next(eval_loader))
872
+ metrics = p_eval_step(state.params, batch)
873
+ eval_metrics.append(metrics)
874
+
875
+ # normalize eval metrics
876
+ eval_metrics = get_metrics(eval_metrics)
877
+ eval_metrics = jax.tree_map(lambda x: jnp.mean(x).item(), eval_metrics)
878
+
879
+ try:
880
+ eval_metrics["perplexity"] = math.exp(eval_metrics["loss"])
881
+ except OverflowError:
882
+ eval_metrics["perplexity"] = float("inf")
883
+
884
+ if jax.process_index() == 0:
885
+ eval_metrics = {f"eval_{metric_name}": value for metric_name, value in eval_metrics.items()}
886
+ path = os.path.join(training_args.output_dir, "eval_results.json")
887
+ with open(path, "w") as f:
888
+ json.dump(eval_metrics, f, indent=4, sort_keys=True)
889
+
890
+
891
+ if __name__ == "__main__":
892
+ main()
special_tokens_map.json ADDED
@@ -0,0 +1 @@
 
 
1
+ {"bos_token": "<|endoftext|>", "eos_token": "<|endoftext|>", "unk_token": "<|endoftext|>", "pad_token": "<|endoftext|>"}
start_train.sh ADDED
@@ -0,0 +1,29 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # set train hyperparams
2
+ unset LD_PRELOAD
3
+ export HF_DATASETS_CACHE="/researchdisk/datasets_cache"
4
+ export USE_TORCH=0
5
+ python3 run_clm_flax.py \
6
+ --output_dir="./" \
7
+ --model_type="gpt2" \
8
+ --config_name="./" \
9
+ --tokenizer_name="./" \
10
+ --dataset_filepath="/researchdisk/training_dataset_full_deduplicated" \
11
+ --do_train --do_eval \
12
+ --block_size="512" \
13
+ --per_device_train_batch_size="16" \
14
+ --per_device_eval_batch_size="16" \
15
+ --preprocessing_num_workers="1" \
16
+ --adam_beta1="0.9" \
17
+ --adam_beta2="0.98" \
18
+ --learning_rate="1e-5" \
19
+ --weight_decay="0.01" \
20
+ --warmup_steps="4000" \
21
+ --cosine_decay \
22
+ --overwrite_output_dir \
23
+ --logging_steps="500" \
24
+ --eval_steps="10000" \
25
+ --save_steps="50000" \
26
+ --num_train_epochs="5" \
27
+ --dtype="bfloat16" \
28
+ --push_to_hub \
29
+ --hub_model_id="Finnish-NLP/gpt2-large-finnish"
tokenizer.json ADDED
The diff for this file is too large to render. See raw diff
 
tokenizer_config.json ADDED
@@ -0,0 +1 @@
 
 
1
+ {"unk_token": "<|endoftext|>", "bos_token": "<|endoftext|>", "eos_token": "<|endoftext|>", "add_prefix_space": false, "special_tokens_map_file": null, "name_or_path": "./", "tokenizer_class": "GPT2Tokenizer"}
train_tokenizer.py ADDED
@@ -0,0 +1,30 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from datasets import load_from_disk
2
+ from tokenizers import trainers, Tokenizer, normalizers, ByteLevelBPETokenizer
3
+ from transformers import AutoTokenizer
4
+
5
+
6
+ model_dir = "./"
7
+
8
+ # load dataset
9
+ dataset = load_from_disk("/researchdisk/training_dataset_full_deduplicated")
10
+ dataset = dataset["train"]
11
+
12
+ # Instantiate tokenizer
13
+ tokenizer = ByteLevelBPETokenizer()
14
+ def batch_iterator(batch_size=1000):
15
+ for i in range(0, len(dataset), batch_size):
16
+ yield dataset[i: i + batch_size]["text"]
17
+
18
+ # Customized training
19
+ tokenizer.train_from_iterator(batch_iterator(), vocab_size=50257, min_frequency=2, special_tokens=[
20
+ "<s>",
21
+ "<pad>",
22
+ "</s>",
23
+ "<unk>",
24
+ "<mask>",
25
+ ])
26
+
27
+ # Save files to disk
28
+ tokenizer.save(f"{model_dir}/tokenizer.json")
29
+ tokenizer = AutoTokenizer.from_pretrained(model_dir)
30
+ tokenizer.save_pretrained(model_dir)
vocab.json ADDED
The diff for this file is too large to render. See raw diff