pere commited on
Commit
1a2e9a5
1 Parent(s): cd5ea82

ready for first training attempt

Browse files
Files changed (2) hide show
  1. run.sh +19 -0
  2. run_clm_flax.py +676 -0
run.sh ADDED
@@ -0,0 +1,19 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ python run_clm_flax.py \
2
+ --output_dir="./" \
3
+ --model_type="gpt2" \
4
+ --config_name="./" \
5
+ --tokenizer_name="./" \
6
+ --train_file="/mnt/disks/flaxdisk/corpus/train.json" \
7
+ --validation_file="/mnt/disks/flaxdisk/corpus/validation.json" \
8
+ --do_train --do_eval \
9
+ --block_size="512" \
10
+ --per_device_train_batch_size="64" \
11
+ --per_device_eval_batch_size="64" \
12
+ --learning_rate="5e-3" --warmup_steps="1000" \
13
+ --adam_beta1="0.9" --adam_beta2="0.98" --weight_decay="0.01" \
14
+ --overwrite_output_dir \
15
+ --num_train_epochs="20" \
16
+ --logging_steps="500" \
17
+ --save_steps="2500" \
18
+ --eval_steps="2500" \
19
+ --push_to_hub
run_clm_flax.py ADDED
@@ -0,0 +1,676 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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=causal-lm
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 logging
25
+ import math
26
+ import os
27
+ import sys
28
+ import time
29
+ from dataclasses import dataclass, field
30
+ from pathlib import Path
31
+ from typing import Callable, Optional
32
+
33
+ import datasets
34
+ import numpy as np
35
+ from datasets import Dataset, load_dataset
36
+ from tqdm import tqdm
37
+
38
+ import jax
39
+ import jax.numpy as jnp
40
+ import optax
41
+ import transformers
42
+ from flax import jax_utils, traverse_util
43
+ from flax.jax_utils import unreplicate
44
+ from flax.training import train_state
45
+ from flax.training.common_utils import get_metrics, onehot, shard, shard_prng_key
46
+ from huggingface_hub import Repository
47
+ from transformers import (
48
+ CONFIG_MAPPING,
49
+ FLAX_MODEL_FOR_CAUSAL_LM_MAPPING,
50
+ AutoConfig,
51
+ AutoTokenizer,
52
+ FlaxAutoModelForCausalLM,
53
+ HfArgumentParser,
54
+ TrainingArguments,
55
+ is_tensorboard_available,
56
+ set_seed,
57
+ )
58
+ from transformers.file_utils import get_full_repo_name
59
+ from transformers.testing_utils import CaptureLogger
60
+
61
+
62
+ logger = logging.getLogger(__name__)
63
+
64
+ MODEL_CONFIG_CLASSES = list(FLAX_MODEL_FOR_CAUSAL_LM_MAPPING.keys())
65
+ MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)
66
+
67
+
68
+ @dataclass
69
+ class ModelArguments:
70
+ """
71
+ Arguments pertaining to which model/config/tokenizer we are going to fine-tune, or train from scratch.
72
+ """
73
+
74
+ model_name_or_path: Optional[str] = field(
75
+ default=None,
76
+ metadata={
77
+ "help": "The model checkpoint for weights initialization."
78
+ "Don't set if you want to train a model from scratch."
79
+ },
80
+ )
81
+ model_type: Optional[str] = field(
82
+ default=None,
83
+ metadata={"help": "If training from scratch, pass a model type from the list: " + ", ".join(MODEL_TYPES)},
84
+ )
85
+ config_name: Optional[str] = field(
86
+ default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
87
+ )
88
+ tokenizer_name: Optional[str] = field(
89
+ default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
90
+ )
91
+ cache_dir: Optional[str] = field(
92
+ default=None, metadata={"help": "Where do you want to store the pretrained models downloaded from s3"}
93
+ )
94
+ use_fast_tokenizer: bool = field(
95
+ default=True,
96
+ metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."},
97
+ )
98
+ dtype: Optional[str] = field(
99
+ default="float32",
100
+ metadata={
101
+ "help": "Floating-point format in which the model weights should be initialized and trained. Choose one of `[float32, float16, bfloat16]`."
102
+ },
103
+ )
104
+
105
+
106
+ @dataclass
107
+ class DataTrainingArguments:
108
+ """
109
+ Arguments pertaining to what data we are going to input our model for training and eval.
110
+ """
111
+
112
+ dataset_name: Optional[str] = field(
113
+ default=None, metadata={"help": "The name of the dataset to use (via the datasets library)."}
114
+ )
115
+ dataset_config_name: Optional[str] = field(
116
+ default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
117
+ )
118
+ train_file: Optional[str] = field(default=None, metadata={"help": "The input training data file (a text file)."})
119
+ validation_file: Optional[str] = field(
120
+ default=None,
121
+ metadata={"help": "An optional input evaluation data file to evaluate the perplexity on (a text file)."},
122
+ )
123
+ max_train_samples: Optional[int] = field(
124
+ default=None,
125
+ metadata={
126
+ "help": "For debugging purposes or quicker training, truncate the number of training examples to this "
127
+ "value if set."
128
+ },
129
+ )
130
+ max_eval_samples: Optional[int] = field(
131
+ default=None,
132
+ metadata={
133
+ "help": "For debugging purposes or quicker training, truncate the number of evaluation examples to this "
134
+ "value if set."
135
+ },
136
+ )
137
+ overwrite_cache: bool = field(
138
+ default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
139
+ )
140
+ validation_split_percentage: Optional[int] = field(
141
+ default=5,
142
+ metadata={
143
+ "help": "The percentage of the train set used as validation set in case there's no validation split"
144
+ },
145
+ )
146
+ block_size: Optional[int] = field(
147
+ default=None,
148
+ metadata={
149
+ "help": "Optional input sequence length after tokenization. "
150
+ "The training dataset will be truncated in block of this size for training. "
151
+ "Default to the model max input length for single sentence inputs (take into account special tokens)."
152
+ },
153
+ )
154
+ overwrite_cache: bool = field(
155
+ default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
156
+ )
157
+ preprocessing_num_workers: Optional[int] = field(
158
+ default=None,
159
+ metadata={"help": "The number of processes to use for the preprocessing."},
160
+ )
161
+ keep_linebreaks: bool = field(
162
+ default=True, metadata={"help": "Whether to keep line breaks when using TXT files or not."}
163
+ )
164
+
165
+ def __post_init__(self):
166
+ if self.dataset_name is None and self.train_file is None and self.validation_file is None:
167
+ raise ValueError("Need either a dataset name or a training/validation file.")
168
+ else:
169
+ if self.train_file is not None:
170
+ extension = self.train_file.split(".")[-1]
171
+ assert extension in ["csv", "json", "txt"], "`train_file` should be a csv, a json or a txt file."
172
+ if self.validation_file is not None:
173
+ extension = self.validation_file.split(".")[-1]
174
+ assert extension in ["csv", "json", "txt"], "`validation_file` should be a csv, a json or a txt file."
175
+
176
+
177
+ class TrainState(train_state.TrainState):
178
+ dropout_rng: jnp.ndarray
179
+
180
+ def replicate(self):
181
+ return jax_utils.replicate(self).replace(dropout_rng=shard_prng_key(self.dropout_rng))
182
+
183
+
184
+ def data_loader(rng: jax.random.PRNGKey, dataset: Dataset, batch_size: int, shuffle: bool = False):
185
+ """
186
+ Returns batches of size `batch_size` from truncated `dataset`, sharded over all local devices.
187
+ Shuffle batches if `shuffle` is `True`.
188
+ """
189
+ steps_per_epoch = len(dataset) // batch_size
190
+
191
+ if shuffle:
192
+ batch_idx = np.random.permutation(len(dataset))
193
+ else:
194
+ batch_idx = np.arange(len(dataset))
195
+
196
+ batch_idx = batch_idx[: steps_per_epoch * batch_size] # Skip incomplete batch.
197
+ batch_idx = batch_idx.reshape((steps_per_epoch, batch_size))
198
+
199
+ for idx in batch_idx:
200
+ batch = dataset[idx]
201
+ batch = {k: np.array(v) for k, v in batch.items()}
202
+
203
+ yield batch
204
+
205
+
206
+ def write_train_metric(summary_writer, train_metrics, train_time, step):
207
+ summary_writer.scalar("train_time", train_time, step)
208
+
209
+ train_metrics = get_metrics(train_metrics)
210
+ for key, vals in train_metrics.items():
211
+ tag = f"train_{key}"
212
+ for i, val in enumerate(vals):
213
+ summary_writer.scalar(tag, val, step - len(vals) + i + 1)
214
+
215
+
216
+ def write_eval_metric(summary_writer, eval_metrics, step):
217
+ for metric_name, value in eval_metrics.items():
218
+ summary_writer.scalar(f"eval_{metric_name}", value, step)
219
+
220
+
221
+ def create_learning_rate_fn(
222
+ train_ds_size: int, train_batch_size: int, num_train_epochs: int, num_warmup_steps: int, learning_rate: float
223
+ ) -> Callable[[int], jnp.array]:
224
+ """Returns a linear warmup, linear_decay learning rate function."""
225
+ steps_per_epoch = train_ds_size // train_batch_size
226
+ num_train_steps = steps_per_epoch * num_train_epochs
227
+ warmup_fn = optax.linear_schedule(init_value=0.0, end_value=learning_rate, transition_steps=num_warmup_steps)
228
+ decay_fn = optax.linear_schedule(
229
+ init_value=learning_rate, end_value=0, transition_steps=num_train_steps - num_warmup_steps
230
+ )
231
+ schedule_fn = optax.join_schedules(schedules=[warmup_fn, decay_fn], boundaries=[num_warmup_steps])
232
+ return schedule_fn
233
+
234
+
235
+ def main():
236
+ # See all possible arguments in src/transformers/training_args.py
237
+ # or by passing the --help flag to this script.
238
+ # We now keep distinct sets of args, for a cleaner separation of concerns.
239
+
240
+ parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
241
+ if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
242
+ # If we pass only one argument to the script and it's the path to a json file,
243
+ # let's parse it to get our arguments.
244
+ model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
245
+ else:
246
+ model_args, data_args, training_args = parser.parse_args_into_dataclasses()
247
+
248
+ if (
249
+ os.path.exists(training_args.output_dir)
250
+ and os.listdir(training_args.output_dir)
251
+ and training_args.do_train
252
+ and not training_args.overwrite_output_dir
253
+ ):
254
+ raise ValueError(
255
+ f"Output directory ({training_args.output_dir}) already exists and is not empty."
256
+ "Use --overwrite_output_dir to overcome."
257
+ )
258
+
259
+ # Make one log on every process with the configuration for debugging.
260
+ logging.basicConfig(
261
+ format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
262
+ datefmt="%m/%d/%Y %H:%M:%S",
263
+ level=logging.INFO,
264
+ )
265
+ # Setup logging, we only want one process per machine to log things on the screen.
266
+ logger.setLevel(logging.INFO if jax.process_index() == 0 else logging.ERROR)
267
+ if jax.process_index() == 0:
268
+ datasets.utils.logging.set_verbosity_warning()
269
+ transformers.utils.logging.set_verbosity_info()
270
+ else:
271
+ datasets.utils.logging.set_verbosity_error()
272
+ transformers.utils.logging.set_verbosity_error()
273
+
274
+ # Set the verbosity to info of the Transformers logger (on main process only):
275
+ logger.info(f"Training/evaluation parameters {training_args}")
276
+
277
+ # Set seed before initializing model.
278
+ set_seed(training_args.seed)
279
+
280
+ # Handle the repository creation
281
+ if training_args.push_to_hub:
282
+ if training_args.hub_model_id is None:
283
+ repo_name = get_full_repo_name(
284
+ Path(training_args.output_dir).absolute().name, token=training_args.hub_token
285
+ )
286
+ else:
287
+ repo_name = training_args.hub_model_id
288
+ repo = Repository(training_args.output_dir, clone_from=repo_name)
289
+
290
+ # Get the datasets: you can either provide your own CSV/JSON/TXT training and evaluation files (see below)
291
+ # or just provide the name of one of the public datasets available on the hub at https://huggingface.co/datasets/
292
+ # (the dataset will be downloaded automatically from the datasets Hub).
293
+ #
294
+ # For CSV/JSON files, this script will use the column called 'text' or the first column if no column called
295
+ # 'text' is found. You can easily tweak this behavior (see below).
296
+ #
297
+ # In distributed training, the load_dataset function guarantees that only one local process can concurrently
298
+ # download the dataset.
299
+ if data_args.dataset_name is not None:
300
+ # Downloading and loading a dataset from the hub.
301
+ dataset = load_dataset(
302
+ data_args.dataset_name, data_args.dataset_config_name, cache_dir=model_args.cache_dir, keep_in_memory=False
303
+ )
304
+
305
+ if "validation" not in dataset.keys():
306
+ dataset["validation"] = load_dataset(
307
+ data_args.dataset_name,
308
+ data_args.dataset_config_name,
309
+ split=f"train[:{data_args.validation_split_percentage}%]",
310
+ cache_dir=model_args.cache_dir,
311
+ )
312
+ dataset["train"] = load_dataset(
313
+ data_args.dataset_name,
314
+ data_args.dataset_config_name,
315
+ split=f"train[{data_args.validation_split_percentage}%:]",
316
+ cache_dir=model_args.cache_dir,
317
+ )
318
+ else:
319
+ data_files = {}
320
+ dataset_args = {}
321
+ if data_args.train_file is not None:
322
+ data_files["train"] = data_args.train_file
323
+ if data_args.validation_file is not None:
324
+ data_files["validation"] = data_args.validation_file
325
+ extension = data_args.train_file.split(".")[-1]
326
+ if extension == "txt":
327
+ extension = "text"
328
+ dataset_args["keep_linebreaks"] = data_args.keep_linebreaks
329
+ dataset = load_dataset(extension, data_files=data_files, cache_dir=model_args.cache_dir, **dataset_args)
330
+
331
+ if "validation" not in dataset.keys():
332
+ dataset["validation"] = load_dataset(
333
+ extension,
334
+ data_files=data_files,
335
+ split=f"train[:{data_args.validation_split_percentage}%]",
336
+ cache_dir=model_args.cache_dir,
337
+ **dataset_args,
338
+ )
339
+ dataset["train"] = load_dataset(
340
+ extension,
341
+ data_files=data_files,
342
+ split=f"train[{data_args.validation_split_percentage}%:]",
343
+ cache_dir=model_args.cache_dir,
344
+ **dataset_args,
345
+ )
346
+ # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at
347
+ # https://huggingface.co/docs/datasets/loading_datasets.html.
348
+
349
+ # Load pretrained model and tokenizer
350
+
351
+ # Distributed training:
352
+ # The .from_pretrained methods guarantee that only one local process can concurrently
353
+ # download model & vocab.
354
+ if model_args.config_name:
355
+ config = AutoConfig.from_pretrained(model_args.config_name, cache_dir=model_args.cache_dir)
356
+ elif model_args.model_name_or_path:
357
+ config = AutoConfig.from_pretrained(model_args.model_name_or_path, cache_dir=model_args.cache_dir)
358
+ else:
359
+ config = CONFIG_MAPPING[model_args.model_type]()
360
+ logger.warning("You are instantiating a new config instance from scratch.")
361
+
362
+ if model_args.tokenizer_name:
363
+ tokenizer = AutoTokenizer.from_pretrained(
364
+ model_args.tokenizer_name, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_tokenizer
365
+ )
366
+ elif model_args.model_name_or_path:
367
+ tokenizer = AutoTokenizer.from_pretrained(
368
+ model_args.model_name_or_path, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_tokenizer
369
+ )
370
+ else:
371
+ raise ValueError(
372
+ "You are instantiating a new tokenizer from scratch. This is not supported by this script."
373
+ "You can do it from another script, save it, and load it from here, using --tokenizer_name."
374
+ )
375
+
376
+ if model_args.model_name_or_path:
377
+ model = FlaxAutoModelForCausalLM.from_pretrained(
378
+ model_args.model_name_or_path, config=config, seed=training_args.seed, dtype=getattr(jnp, model_args.dtype)
379
+ )
380
+ else:
381
+ model = FlaxAutoModelForCausalLM.from_config(
382
+ config, seed=training_args.seed, dtype=getattr(jnp, model_args.dtype)
383
+ )
384
+
385
+ # Preprocessing the datasets.
386
+ # First we tokenize all the texts.
387
+ if training_args.do_train:
388
+ column_names = dataset["train"].column_names
389
+ else:
390
+ column_names = dataset["validation"].column_names
391
+ text_column_name = "text" if "text" in column_names else column_names[0]
392
+
393
+ # since this will be pickled to avoid _LazyModule error in Hasher force logger loading before tokenize_function
394
+ tok_logger = transformers.utils.logging.get_logger("transformers.tokenization_utils_base")
395
+
396
+ def tokenize_function(examples):
397
+ with CaptureLogger(tok_logger) as cl:
398
+ output = tokenizer(examples[text_column_name])
399
+ # clm input could be much much longer than block_size
400
+ if "Token indices sequence length is longer than the" in cl.out:
401
+ tok_logger.warning(
402
+ "^^^^^^^^^^^^^^^^ Please ignore the warning above - this long input will be chunked into smaller bits before being passed to the model."
403
+ )
404
+ return output
405
+
406
+ tokenized_datasets = dataset.map(
407
+ tokenize_function,
408
+ batched=True,
409
+ num_proc=data_args.preprocessing_num_workers,
410
+ remove_columns=column_names,
411
+ load_from_cache_file=not data_args.overwrite_cache,
412
+ )
413
+
414
+ if data_args.block_size is None:
415
+ block_size = tokenizer.model_max_length
416
+ if block_size > config.max_position_embeddings:
417
+ logger.warning(
418
+ f"The tokenizer picked seems to have a very large `model_max_length` ({tokenizer.model_max_length}). "
419
+ "Picking 1024 instead. You can change that default value by passing --block_size xxx."
420
+ )
421
+ block_size = 1024
422
+ else:
423
+ if data_args.block_size > tokenizer.model_max_length:
424
+ logger.warning(
425
+ f"The block_size passed ({data_args.block_size}) is larger than the maximum length for the model"
426
+ f"({tokenizer.model_max_length}). Using block_size={tokenizer.model_max_length}."
427
+ )
428
+ block_size = min(data_args.block_size, tokenizer.model_max_length)
429
+
430
+ # Main data processing function that will concatenate all texts from our dataset and generate chunks of block_size.
431
+ def group_texts(examples):
432
+ # Concatenate all texts.
433
+ concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}
434
+ total_length = len(concatenated_examples[list(examples.keys())[0]])
435
+ # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can
436
+ # customize this part to your needs.
437
+ if total_length >= block_size:
438
+ total_length = (total_length // block_size) * block_size
439
+ # Split by chunks of max_len.
440
+ result = {
441
+ k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
442
+ for k, t in concatenated_examples.items()
443
+ }
444
+ result["labels"] = result["input_ids"].copy()
445
+ return result
446
+
447
+ # Note that with `batched=True`, this map processes 1,000 texts together, so group_texts throws away a remainder
448
+ # for each of those groups of 1,000 texts. You can adjust that batch_size here but a higher value might be slower
449
+ # to preprocess.
450
+ #
451
+ # To speed up this part, we use multiprocessing. See the documentation of the map method for more information:
452
+ # https://huggingface.co/docs/datasets/package_reference/main_classes.html#datasets.Dataset.map
453
+
454
+ lm_datasets = tokenized_datasets.map(
455
+ group_texts,
456
+ batched=True,
457
+ num_proc=data_args.preprocessing_num_workers,
458
+ load_from_cache_file=not data_args.overwrite_cache,
459
+ )
460
+
461
+ if training_args.do_train:
462
+ if "train" not in tokenized_datasets:
463
+ raise ValueError("--do_train requires a train dataset")
464
+ train_dataset = lm_datasets["train"]
465
+ if data_args.max_train_samples is not None:
466
+ train_dataset = train_dataset.select(range(data_args.max_train_samples))
467
+
468
+ if training_args.do_eval:
469
+ if "validation" not in tokenized_datasets:
470
+ raise ValueError("--do_eval requires a validation dataset")
471
+ eval_dataset = lm_datasets["validation"]
472
+ if data_args.max_eval_samples is not None:
473
+ eval_dataset = eval_dataset.select(range(data_args.max_eval_samples))
474
+
475
+ # Enable tensorboard only on the master node
476
+ has_tensorboard = is_tensorboard_available()
477
+ if has_tensorboard and jax.process_index() == 0:
478
+ try:
479
+ from flax.metrics.tensorboard import SummaryWriter
480
+
481
+ summary_writer = SummaryWriter(log_dir=Path(training_args.output_dir))
482
+ except ImportError as ie:
483
+ has_tensorboard = False
484
+ logger.warning(
485
+ f"Unable to display metrics through TensorBoard because some package are not installed: {ie}"
486
+ )
487
+ else:
488
+ logger.warning(
489
+ "Unable to display metrics through TensorBoard because the package is not installed: "
490
+ "Please run pip install tensorboard to enable."
491
+ )
492
+
493
+ # Initialize our training
494
+ rng = jax.random.PRNGKey(training_args.seed)
495
+ rng, dropout_rng = jax.random.split(rng)
496
+
497
+ # Store some constant
498
+ num_epochs = int(training_args.num_train_epochs)
499
+ train_batch_size = int(training_args.per_device_train_batch_size) * jax.device_count()
500
+ eval_batch_size = int(training_args.per_device_eval_batch_size) * jax.device_count()
501
+ steps_per_epoch = len(train_dataset) // train_batch_size
502
+ total_train_steps = steps_per_epoch * num_epochs
503
+
504
+ # Create learning rate schedule
505
+ linear_decay_lr_schedule_fn = create_learning_rate_fn(
506
+ len(train_dataset),
507
+ train_batch_size,
508
+ training_args.num_train_epochs,
509
+ training_args.warmup_steps,
510
+ training_args.learning_rate,
511
+ )
512
+
513
+ # We use Optax's "masking" functionality to not apply weight decay
514
+ # to bias and LayerNorm scale parameters. decay_mask_fn returns a
515
+ # mask boolean with the same structure as the parameters.
516
+ # The mask is True for parameters that should be decayed.
517
+ # Note that this mask is specifically adapted for FlaxGPT2.
518
+ # For other models, one should correct the layer norm parameter naming
519
+ # accordingly.
520
+ def decay_mask_fn(params):
521
+ flat_params = traverse_util.flatten_dict(params)
522
+ flat_mask = {
523
+ path: (path[-1] != "bias" and path[-2:] not in [("ln_1", "scale"), ("ln_2", "scale"), ("ln_f", "scale")])
524
+ for path in flat_params
525
+ }
526
+ return traverse_util.unflatten_dict(flat_mask)
527
+
528
+ # create adam optimizer
529
+ if training_args.adafactor:
530
+ # We use the default parameters here to initialize adafactor,
531
+ # For more details about the parameters please check https://github.com/deepmind/optax/blob/ed02befef9bf81cbbf236be3d2b0e032e9ed4a40/optax/_src/alias.py#L74
532
+ optimizer = optax.adafactor(
533
+ learning_rate=linear_decay_lr_schedule_fn,
534
+ )
535
+ else:
536
+ optimizer = optax.adamw(
537
+ learning_rate=linear_decay_lr_schedule_fn,
538
+ b1=training_args.adam_beta1,
539
+ b2=training_args.adam_beta2,
540
+ eps=training_args.adam_epsilon,
541
+ weight_decay=training_args.weight_decay,
542
+ mask=decay_mask_fn,
543
+ )
544
+
545
+ # Setup train state
546
+ state = TrainState.create(apply_fn=model.__call__, params=model.params, tx=optimizer, dropout_rng=dropout_rng)
547
+
548
+ def loss_fn(logits, labels):
549
+ shift_logits = logits[..., :-1, :]
550
+ shift_labels = labels[..., 1:]
551
+ loss = optax.softmax_cross_entropy(shift_logits, onehot(shift_labels, shift_logits.shape[-1]))
552
+ return loss.mean()
553
+
554
+ # Define gradient update step fn
555
+ def train_step(state, batch):
556
+ dropout_rng, new_dropout_rng = jax.random.split(state.dropout_rng)
557
+
558
+ def compute_loss(params):
559
+ labels = batch.pop("labels")
560
+ logits = state.apply_fn(**batch, params=params, dropout_rng=dropout_rng, train=True)[0]
561
+ loss = loss_fn(logits, labels)
562
+ return loss
563
+
564
+ grad_fn = jax.value_and_grad(compute_loss)
565
+ loss, grad = grad_fn(state.params)
566
+ grad = jax.lax.pmean(grad, "batch")
567
+
568
+ new_state = state.apply_gradients(grads=grad, dropout_rng=new_dropout_rng)
569
+
570
+ metrics = {"loss": loss, "learning_rate": linear_decay_lr_schedule_fn(state.step)}
571
+ metrics = jax.lax.pmean(metrics, axis_name="batch")
572
+
573
+ return new_state, metrics
574
+
575
+ # Define eval fn
576
+ def eval_step(params, batch):
577
+ labels = batch.pop("labels")
578
+ logits = model(**batch, params=params, train=False)[0]
579
+ loss = loss_fn(logits, labels)
580
+
581
+ # summarize metrics
582
+ metrics = {"loss": loss}
583
+ metrics = jax.lax.pmean(metrics, axis_name="batch")
584
+ return metrics
585
+
586
+ # Create parallel version of the train and eval step
587
+ p_train_step = jax.pmap(train_step, "batch", donate_argnums=(0,))
588
+ p_eval_step = jax.pmap(eval_step, "batch")
589
+
590
+ # Replicate the train state on each device
591
+ state = state.replicate()
592
+
593
+ logger.info("***** Running training *****")
594
+ logger.info(f" Num examples = {len(train_dataset)}")
595
+ logger.info(f" Num Epochs = {num_epochs}")
596
+ logger.info(f" Instantaneous batch size per device = {training_args.per_device_train_batch_size}")
597
+ logger.info(f" Total train batch size (w. parallel & distributed) = {train_batch_size}")
598
+ logger.info(f" Total optimization steps = {total_train_steps}")
599
+
600
+ train_time = 0
601
+ train_metrics = []
602
+ epochs = tqdm(range(num_epochs), desc="Epoch ... ", position=0)
603
+ for epoch in epochs:
604
+ # ======================== Training ================================
605
+ train_start = time.time()
606
+
607
+ # Create sampling rng
608
+ rng, input_rng = jax.random.split(rng)
609
+
610
+ # Generate an epoch by shuffling sampling indices from the train dataset
611
+ train_loader = data_loader(input_rng, train_dataset, train_batch_size, shuffle=True)
612
+ steps_per_epoch = len(train_dataset) // train_batch_size
613
+ # train
614
+ for step in tqdm(range(steps_per_epoch), desc="Training...", position=1, leave=False):
615
+ batch = next(train_loader)
616
+ batch = shard(batch)
617
+ state, train_metric = p_train_step(state, batch)
618
+ train_metrics.append(train_metric)
619
+
620
+ cur_step = epoch * (len(train_dataset) // train_batch_size) + step
621
+
622
+ if cur_step % training_args.logging_steps == 0 and cur_step > 0:
623
+ # Save metrics
624
+ train_metric = unreplicate(train_metric)
625
+ train_time += time.time() - train_start
626
+ if has_tensorboard and jax.process_index() == 0:
627
+ write_train_metric(summary_writer, train_metrics, train_time, cur_step)
628
+
629
+ epochs.write(
630
+ f"Step... ({cur_step} | Loss: {train_metric['loss'].mean()}, Learning Rate: {train_metric['learning_rate'].mean()})"
631
+ )
632
+
633
+ train_metrics = []
634
+
635
+ if cur_step % training_args.eval_steps == 0 and cur_step > 0:
636
+ # ======================== Evaluating ==============================
637
+ eval_metrics = []
638
+ eval_loader = data_loader(input_rng, eval_dataset, eval_batch_size)
639
+ eval_steps = len(eval_dataset) // eval_batch_size
640
+ for _ in tqdm(range(eval_steps), desc="Evaluating...", position=2, leave=False):
641
+ # Model forward
642
+ batch = next(eval_loader)
643
+ batch = shard(batch)
644
+ metrics = p_eval_step(state.params, batch)
645
+ eval_metrics.append(metrics)
646
+
647
+ # normalize eval metrics
648
+ eval_metrics = get_metrics(eval_metrics)
649
+ eval_metrics = jax.tree_map(jnp.mean, eval_metrics)
650
+
651
+ try:
652
+ eval_metrics["perplexity"] = math.exp(eval_metrics["loss"])
653
+ except OverflowError:
654
+ eval_metrics["perplexity"] = float("inf")
655
+
656
+ # Print metrics and update progress bar
657
+ desc = f"Step... ({cur_step} | Eval Loss: {eval_metrics['loss']} | Eval Perplexity: {eval_metrics['perplexity']})"
658
+ epochs.write(desc)
659
+ epochs.desc = desc
660
+
661
+ # Save metrics
662
+ if has_tensorboard and jax.process_index() == 0:
663
+ write_eval_metric(summary_writer, eval_metrics, cur_step)
664
+
665
+ if cur_step % training_args.save_steps == 0 and cur_step > 0:
666
+ # save checkpoint after each epoch and push checkpoint to the hub
667
+ if jax.process_index() == 0:
668
+ params = jax.device_get(unreplicate(state.params))
669
+ model.save_pretrained(training_args.output_dir, params=params)
670
+ tokenizer.save_pretrained(training_args.output_dir)
671
+ if training_args.push_to_hub:
672
+ repo.push_to_hub(commit_message=f"Saving weights and logs of step {cur_step}", blocking=False)
673
+
674
+
675
+ if __name__ == "__main__":
676
+ main()