yhavinga commited on
Commit
c6e3db9
1 Parent(s): 055f129

Add scripts, model config etc

Browse files
.gitignore ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ data
2
+ *~
README.md ADDED
@@ -0,0 +1,44 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ language: nl
3
+ widget:
4
+ - text: "In het jaar 2030 zullen we"
5
+ - text: "Toen ik gisteren volledig in de ban was van"
6
+ - text: "Studenten en leraren van de Bogazici Universiteit in de Turkse stad Istanbul"
7
+ - text: "In Israël was een strenge lockdown"
8
+ tags:
9
+ - gpt2-large
10
+ - gpt2
11
+ pipeline_tag: text-generation
12
+ datasets:
13
+ - yhavinga/mc4_nl_cleaned
14
+ ---
15
+ # GPT2-Large pre-trained on cleaned Dutch mC4 🇳🇱
16
+
17
+ Dataset:
18
+
19
+ * [mC4 NL Cleaned](https://huggingface.co/datasets/yhavinga/mc4_nl_cleaned)
20
+ * dataset config: full (33B tokens)
21
+
22
+ Tokenizer:
23
+
24
+ * Tokenizer trained on mC4 with scripts from the Huggingface
25
+ Transformers [Flax examples](https://github.com/huggingface/transformers/tree/master/examples/flax/language-modeling)
26
+
27
+ Training details:
28
+
29
+ * Training at step 360000 of 2082009 (17%)
30
+ * Block size: 512
31
+ * Optimizer: adafactor
32
+ * Learning rate: 3.3e-5
33
+ * Batch size: 32
34
+ * Warmup steps: 5000
35
+ * Weight decay: 0.01
36
+
37
+ Work in progress. Dec 2021-Jan2022
38
+
39
+ * Many thanks to the [Google TPU Research Cloud](https://sites.research.google/trc/about/) for providing access to a TPU cluster!
40
+ * Thanks to @gsarti for creating the [t5-flax-gcp
41
+ repository](https://github.com/gsarti/t5-flax-gcp).
42
+ * Also thanks to the creators of [gpt2-medium-persian](https://huggingface.co/flax-community/gpt2-medium-persian) and
43
+ [gpt2-medium-indonesian](https://huggingface.co/flax-community/gpt2-medium-persian)
44
+ for sharing their training scripts!
added_tokens.json ADDED
@@ -0,0 +1 @@
 
 
1
+ {"<|endoftext|>": 50256}
commit.sh ADDED
@@ -0,0 +1,21 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/bin/bash
2
+ source ~/venv/bin/activate
3
+
4
+ while true
5
+ do
6
+ echo -n "Checking at .. "
7
+ date
8
+ UPDATED=`git status | grep flax_model | grep modified`
9
+
10
+ if [ ! -z "$UPDATED" ]
11
+ then
12
+ sleep 120
13
+ FILE=$(find . -name `ls -tR runs | grep events | head -n 1` | tail -n 1)
14
+ STEP=`tensorboard --load_fast=true --inspect --event_file=$FILE | grep last_step | awk '{print $2}'`
15
+ git add runs
16
+ git add flax_model.msgpack
17
+ git commit -m "Saving weights and logs step $STEP"
18
+ fi
19
+
20
+ sleep 60
21
+ done
config.json ADDED
@@ -0,0 +1,37 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "activation_function": "gelu_new",
3
+ "architectures": [
4
+ "GPT2LMHeadModel"
5
+ ],
6
+ "attn_pdrop": 0.0,
7
+ "bos_token_id": 50256,
8
+ "embd_pdrop": 0.0,
9
+ "eos_token_id": 50256,
10
+ "initializer_range": 0.02,
11
+ "layer_norm_epsilon": 1e-05,
12
+ "model_type": "gpt2",
13
+ "n_ctx": 1024,
14
+ "n_embd": 1280,
15
+ "n_head": 20,
16
+ "n_inner": null,
17
+ "n_layer": 36,
18
+ "n_positions": 1024,
19
+ "reorder_and_upcast_attn": false,
20
+ "resid_pdrop": 0.0,
21
+ "scale_attn_by_inverse_layer_idx": false,
22
+ "scale_attn_weights": true,
23
+ "summary_activation": null,
24
+ "summary_first_dropout": 0.1,
25
+ "summary_proj_to_labels": true,
26
+ "summary_type": "cls_index",
27
+ "summary_use_proj": true,
28
+ "task_specific_params": {
29
+ "text-generation": {
30
+ "do_sample": true,
31
+ "max_length": 50
32
+ }
33
+ },
34
+ "transformers_version": "4.13.0",
35
+ "use_cache": true,
36
+ "vocab_size": 50257
37
+ }
create_config.py ADDED
@@ -0,0 +1,14 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import GPT2Config
2
+
3
+ import os
4
+
5
+ config_type = os.environ.get("CONFIG_TYPE")
6
+ dataset_name = os.environ.get("DATASET")
7
+ dataset_config = os.environ.get("DATASET_CONFIG")
8
+ dataset_split = os.environ.get("DATASET_SPLIT")
9
+ vocab_size = int(os.environ.get("VOCAB_SIZE"))
10
+ model_path = os.environ.get("MODEL_PATH")
11
+
12
+
13
+ config = GPT2Config.from_pretrained(config_type, resid_pdrop=0.0, embd_pdrop=0.0, attn_pdrop=0.0, vocab_size=vocab_size)
14
+ config.save_pretrained(model_path)
flax_to_pytorch.py ADDED
@@ -0,0 +1,22 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import numpy as np
3
+ import jax
4
+ import jax.numpy as jnp
5
+ from transformers import AutoTokenizer
6
+ from transformers import FlaxGPT2LMHeadModel
7
+ from transformers import GPT2LMHeadModel
8
+ tokenizer = AutoTokenizer.from_pretrained(".")
9
+ tokenizer.pad_token = tokenizer.eos_token
10
+ model_fx = FlaxGPT2LMHeadModel.from_pretrained(".")
11
+ # def to_f32(t):
12
+ # return jax.tree_map(lambda x: x.astype(jnp.float32) if x.dtype == jnp.bfloat16 else x, t)
13
+ # model_fx.params = to_f32(model_fx.params)
14
+ # model_fx.save_pretrained("./fx")
15
+ model_pt = GPT2LMHeadModel.from_pretrained(".", from_flax=True)
16
+ model_pt.save_pretrained(".")
17
+ input_ids = np.asarray(2 * [128 * [0]], dtype=np.int32)
18
+ input_ids_pt = torch.tensor(input_ids)
19
+ logits_pt = model_pt(input_ids_pt).logits
20
+ print(logits_pt)
21
+ logits_fx = model_fx(input_ids).logits
22
+ print(logits_fx)
merges.txt ADDED
The diff for this file is too large to render. See raw diff
 
push.sh ADDED
@@ -0,0 +1,16 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/bin/bash
2
+ source ~/venv/bin/activate
3
+
4
+ while true
5
+ do
6
+ echo -n "Checking at .. "
7
+ date
8
+ BEHIND=`git rev-list origin..HEAD`
9
+
10
+ if [ ! -z "$BEHIND" ]
11
+ then
12
+ git push origin
13
+ fi
14
+
15
+ sleep 180
16
+ done
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,895 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+ from dataclasses import asdict, dataclass, field
31
+ from enum import Enum
32
+ from itertools import chain
33
+ from pathlib import Path
34
+ from typing import Callable, Optional
35
+ import json
36
+ import shutil
37
+
38
+ import datasets
39
+ import numpy as np
40
+ from datasets import Dataset, load_dataset
41
+ from tqdm import tqdm
42
+
43
+ import jax
44
+ import jax.numpy as jnp
45
+ import optax
46
+ import transformers
47
+ from flax import jax_utils, traverse_util
48
+ from flax.jax_utils import unreplicate
49
+ from flax.training import train_state
50
+ # from flax.training.checkpoints import save_checkpoint, restore_checkpoint
51
+ from flax.training.common_utils import get_metrics, onehot, shard, shard_prng_key
52
+ from flax.serialization import to_bytes, from_bytes
53
+ from transformers import (
54
+ CONFIG_MAPPING,
55
+ FLAX_MODEL_FOR_CAUSAL_LM_MAPPING,
56
+ AutoConfig,
57
+ AutoTokenizer,
58
+ FlaxAutoModelForCausalLM,
59
+ HfArgumentParser,
60
+ is_tensorboard_available,
61
+ set_seed,
62
+ )
63
+ from transformers.file_utils import get_full_repo_name
64
+ from transformers.testing_utils import CaptureLogger
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
+ learning_rate: float = field(default=5e-5, metadata={"help": "The initial learning rate for AdamW."})
96
+ weight_decay: float = field(default=0.0, metadata={"help": "Weight decay for AdamW if we apply some."})
97
+ adam_beta1: float = field(default=0.9, metadata={"help": "Beta1 for AdamW optimizer"})
98
+ adam_beta2: float = field(default=0.999, metadata={"help": "Beta2 for AdamW optimizer"})
99
+ adam_epsilon: float = field(default=1e-8, metadata={"help": "Epsilon for AdamW optimizer."})
100
+ adafactor: bool = field(default=False, metadata={"help": "Whether or not to replace AdamW by Adafactor."})
101
+ num_train_epochs: float = field(default=3.0, metadata={"help": "Total number of training epochs to perform."})
102
+ warmup_steps: int = field(default=0, metadata={"help": "Linear warmup over warmup_steps."})
103
+ logging_steps: int = field(default=500, metadata={"help": "Log every X updates steps."})
104
+ save_steps: int = field(default=500, metadata={"help": "Save checkpoint every X updates steps."})
105
+ eval_steps: int = field(default=None, metadata={"help": "Run an evaluation every X steps."})
106
+ resume_from_checkpoint: Optional[str] = field(
107
+ default=None,
108
+ metadata={
109
+ "help": "The model checkpoint to resume training from. Should contain training state"
110
+ },
111
+ )
112
+ seed: int = field(default=42, metadata={"help": "Random seed that will be set at the beginning of training."})
113
+ push_to_hub: bool = field(
114
+ default=False, metadata={"help": "Whether or not to upload the trained model to the model hub after training."}
115
+ )
116
+ hub_model_id: str = field(
117
+ default=None, metadata={"help": "The name of the repository to keep in sync with the local `output_dir`."}
118
+ )
119
+ hub_token: str = field(default=None, metadata={"help": "The token to use to push to the Model Hub."})
120
+
121
+ def __post_init__(self):
122
+ if self.output_dir is not None:
123
+ self.output_dir = os.path.expanduser(self.output_dir)
124
+
125
+ def to_dict(self):
126
+ """
127
+ Serializes this instance while replace `Enum` by their values (for JSON serialization support). It obfuscates
128
+ the token values by removing their value.
129
+ """
130
+ d = asdict(self)
131
+ for k, v in d.items():
132
+ if isinstance(v, Enum):
133
+ d[k] = v.value
134
+ if isinstance(v, list) and len(v) > 0 and isinstance(v[0], Enum):
135
+ d[k] = [x.value for x in v]
136
+ if k.endswith("_token"):
137
+ d[k] = f"<{k.upper()}>"
138
+ return d
139
+
140
+
141
+ @dataclass
142
+ class ModelArguments:
143
+ """
144
+ Arguments pertaining to which model/config/tokenizer we are going to fine-tune, or train from scratch.
145
+ """
146
+
147
+ model_name_or_path: Optional[str] = field(
148
+ default=None,
149
+ metadata={
150
+ "help": "The model checkpoint for weights initialization."
151
+ "Don't set if you want to train a model from scratch."
152
+ },
153
+ )
154
+ model_type: Optional[str] = field(
155
+ default=None,
156
+ metadata={"help": "If training from scratch, pass a model type from the list: " + ", ".join(MODEL_TYPES)},
157
+ )
158
+ config_name: Optional[str] = field(
159
+ default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
160
+ )
161
+ tokenizer_name: Optional[str] = field(
162
+ default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
163
+ )
164
+ cache_dir: Optional[str] = field(
165
+ default=None, metadata={"help": "Where do you want to store the pretrained models downloaded from s3"}
166
+ )
167
+ use_fast_tokenizer: bool = field(
168
+ default=True,
169
+ metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."},
170
+ )
171
+ dtype: Optional[str] = field(
172
+ default="float32",
173
+ metadata={
174
+ "help": "Floating-point format in which the model weights should be initialized and trained. Choose one of `[float32, float16, bfloat16]`."
175
+ },
176
+ )
177
+
178
+
179
+ @dataclass
180
+ class DataTrainingArguments:
181
+ """
182
+ Arguments pertaining to what data we are going to input our model for training and eval.
183
+ """
184
+
185
+ dataset_name: Optional[str] = field(
186
+ default=None, metadata={"help": "The name of the dataset to use (via the datasets library)."}
187
+ )
188
+ dataset_config_name: Optional[str] = field(
189
+ default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
190
+ )
191
+ train_file: Optional[str] = field(default=None, metadata={"help": "The input training data file (a text file)."})
192
+ validation_file: Optional[str] = field(
193
+ default=None,
194
+ metadata={"help": "An optional input evaluation data file to evaluate the perplexity on (a text file)."},
195
+ )
196
+ max_train_samples: Optional[int] = field(
197
+ default=None,
198
+ metadata={
199
+ "help": "For debugging purposes or quicker training, truncate the number of training examples to this "
200
+ "value if set."
201
+ },
202
+ )
203
+ max_eval_samples: Optional[int] = field(
204
+ default=None,
205
+ metadata={
206
+ "help": "For debugging purposes or quicker training, truncate the number of evaluation examples to this "
207
+ "value if set."
208
+ },
209
+ )
210
+ overwrite_cache: bool = field(
211
+ default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
212
+ )
213
+ validation_split_percentage: Optional[int] = field(
214
+ default=5,
215
+ metadata={
216
+ "help": "The percentage of the train set used as validation set in case there's no validation split"
217
+ },
218
+ )
219
+ block_size: Optional[int] = field(
220
+ default=None,
221
+ metadata={
222
+ "help": "Optional input sequence length after tokenization. "
223
+ "The training dataset will be truncated in block of this size for training. "
224
+ "Default to the model max input length for single sentence inputs (take into account special tokens)."
225
+ },
226
+ )
227
+ overwrite_cache: bool = field(
228
+ default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
229
+ )
230
+ preprocessing_num_workers: Optional[int] = field(
231
+ default=None,
232
+ metadata={"help": "The number of processes to use for the preprocessing."},
233
+ )
234
+ keep_linebreaks: bool = field(
235
+ default=True, metadata={"help": "Whether to keep line breaks when using TXT files or not."}
236
+ )
237
+
238
+ def __post_init__(self):
239
+ if self.dataset_name is None and self.train_file is None and self.validation_file is None:
240
+ raise ValueError("Need either a dataset name or a training/validation file.")
241
+ else:
242
+ if self.train_file is not None:
243
+ extension = self.train_file.split(".")[-1]
244
+ assert extension in ["csv", "json", "txt"], "`train_file` should be a csv, a json or a txt file."
245
+ if self.validation_file is not None:
246
+ extension = self.validation_file.split(".")[-1]
247
+ assert extension in ["csv", "json", "txt"], "`validation_file` should be a csv, a json or a txt file."
248
+
249
+
250
+ class TrainState(train_state.TrainState):
251
+ dropout_rng: jnp.ndarray
252
+
253
+ def replicate(self):
254
+ return jax_utils.replicate(self).replace(dropout_rng=shard_prng_key(self.dropout_rng))
255
+
256
+
257
+ def data_loader(rng: jax.random.PRNGKey, dataset: Dataset, batch_size: int, shuffle: bool = False):
258
+ """
259
+ Returns batches of size `batch_size` from truncated `dataset`, sharded over all local devices.
260
+ Shuffle batches if `shuffle` is `True`.
261
+ """
262
+ steps_per_epoch = len(dataset) // batch_size
263
+
264
+ if shuffle:
265
+ batch_idx = jax.random.permutation(rng, len(dataset))
266
+ else:
267
+ batch_idx = jnp.arange(len(dataset))
268
+
269
+ batch_idx = batch_idx[: steps_per_epoch * batch_size] # Skip incomplete batch.
270
+ batch_idx = batch_idx.reshape((steps_per_epoch, batch_size))
271
+
272
+ for idx in batch_idx:
273
+ batch = dataset[idx]
274
+ batch = {k: np.array(v) for k, v in batch.items()}
275
+
276
+ yield batch
277
+
278
+
279
+ def write_train_metric(summary_writer, train_metrics, train_time, step):
280
+ summary_writer.scalar("train_time", train_time, step)
281
+
282
+ train_metrics = get_metrics(train_metrics)
283
+ for key, vals in train_metrics.items():
284
+ tag = f"train_{key}"
285
+ for i, val in enumerate(vals):
286
+ summary_writer.scalar(tag, val, step - len(vals) + i + 1)
287
+
288
+
289
+ def write_eval_metric(summary_writer, eval_metrics, step):
290
+ for metric_name, value in eval_metrics.items():
291
+ summary_writer.scalar(f"eval_{metric_name}", value, step)
292
+
293
+
294
+ def create_learning_rate_fn(
295
+ train_ds_size: int, train_batch_size: int, num_train_epochs: int, num_warmup_steps: int, learning_rate: float
296
+ ) -> Callable[[int], jnp.array]:
297
+ """Returns a linear warmup, linear_decay learning rate function."""
298
+ steps_per_epoch = train_ds_size // train_batch_size
299
+ num_train_steps = steps_per_epoch * num_train_epochs
300
+ warmup_fn = optax.linear_schedule(init_value=0.0, end_value=learning_rate, transition_steps=num_warmup_steps)
301
+ decay_fn = optax.linear_schedule(
302
+ init_value=learning_rate, end_value=0, transition_steps=num_train_steps - num_warmup_steps
303
+ )
304
+ schedule_fn = optax.join_schedules(schedules=[warmup_fn, decay_fn], boundaries=[num_warmup_steps])
305
+ return schedule_fn
306
+
307
+
308
+ # utils
309
+ def mb_item(x):
310
+ return x.item() if hasattr(x, "item") else x
311
+
312
+
313
+ # checkpoint functions
314
+ def save_model_checkpoint(model, save_dir, state, with_opt: bool = True, push_to_hub: bool = False):
315
+ """
316
+ If `push_to_hub` is True, will save to `save_dir`. Otherwise will save to `save_dir/ckpt-{step}`.
317
+ """
318
+ state = jax_utils.unreplicate(state)
319
+ logger.info(f"SAVING CHECKPOINT IN {save_dir}...")
320
+ if not push_to_hub:
321
+ save_dir = f"{save_dir}/ckpt-{mb_item(state.step) - 1}"
322
+ model.save_pretrained(
323
+ save_dir,
324
+ params=state.params,
325
+ push_to_hub=push_to_hub,
326
+ commit_message=f"Saving weights and logs at step {mb_item(state.step) - 1}",
327
+ )
328
+ if with_opt:
329
+ with open(os.path.join(save_dir, "opt_state.msgpack"), "wb") as f:
330
+ f.write(to_bytes(state.opt_state))
331
+ with open(os.path.join(save_dir, "training_state.json"), "w") as f:
332
+ json.dump({"step": state.step.item()}, f)
333
+ logger.info("checkpoint saved")
334
+
335
+
336
+ # this is added to make resuming from checkpoint to work with adafactor
337
+ # to be removed when issue is fixed
338
+ # notice that adafactor state is perturbed by fake_update
339
+ def _zeros_tree_like(inp_tree):
340
+ return jax.tree_map(jnp.zeros_like, inp_tree)
341
+
342
+
343
+ def fake_update(state):
344
+ fake_updates = _zeros_tree_like(state.params)
345
+ _, new_inner_opt_state = state.tx.inner_opt.update(fake_updates, state.opt_state.inner_opt_state, state.params)
346
+ opt_state = state.opt_state
347
+ new_opt_state = optax.MultiStepsState(mini_step=opt_state.mini_step,
348
+ gradient_step=opt_state.gradient_step,
349
+ inner_opt_state=new_inner_opt_state,
350
+ acc_grads=opt_state.acc_grads)
351
+ return state.replace(opt_state=new_opt_state)
352
+
353
+
354
+ def reinstantiate_states(opt_state):
355
+ new_state = []
356
+ for state in opt_state:
357
+ if isinstance(state, list):
358
+ new_state.append(reinstantiate_states(state))
359
+ else:
360
+ cls = getattr(optax, type(state).__name__)
361
+ new_state.append(cls(**{k: getattr(state, k) for k in state._fields}))
362
+ return new_state
363
+
364
+
365
+ def restore_model_checkpoint(save_dir, state):
366
+ logger.info(f"RESTORING CHECKPOINT FROM {save_dir}...")
367
+ with open(os.path.join(save_dir, "flax_model.msgpack"), "rb") as f:
368
+ params = from_bytes(state.params, f.read())
369
+
370
+ with open(os.path.join(save_dir, "opt_state.msgpack"), "rb") as f:
371
+ opt_state = from_bytes(state.opt_state, f.read())
372
+
373
+ with open(os.path.join(save_dir, "training_state.json"), "r") as f:
374
+ training_state = json.load(f)
375
+ step = training_state["step"]
376
+
377
+ logger.info("checkpoint restored")
378
+ # reinstantiate inner opt state to avoid type conflict
379
+ if hasattr(opt_state, "inner_opt_state"):
380
+ print("restoring state of multisteps optimizer")
381
+ inner_opt_state = reinstantiate_states(opt_state.inner_opt_state)
382
+ ms_state_dict = {k: getattr(state.opt_state, k) for k in state.opt_state._fields}
383
+ ms_state_dict["inner_opt_state"] = inner_opt_state
384
+ opt_state = optax.MultiStepsState(**ms_state_dict)
385
+
386
+ return state.replace(step=step, params=params, opt_state=opt_state)
387
+
388
+
389
+ def rotate_checkpoints(ckpt_dir: str, save_total_limit: int):
390
+ "Removes older checkpoints so that `save_total_limit` checkpoints are kept"
391
+ # TODO: what to remove is decided using step number only, we might want to improve that
392
+ ckpts = [str(x) for x in Path(ckpt_dir).glob("ckpt-*")]
393
+ # sort checkpoints by step
394
+ ckpts_sorted = sorted(ckpts, key=lambda x: int(x.split('-')[-1]))
395
+ ckpts_to_delete = ckpts_sorted[:-save_total_limit]
396
+ for ckpt in ckpts_to_delete:
397
+ logger.info(f"Deleting older checkpoint [{ckpt}] due to save_total_limit ({save_total_limit})")
398
+ shutil.rmtree(ckpt)
399
+
400
+
401
+ def main():
402
+ # See all possible arguments in src/transformers/training_args.py
403
+ # or by passing the --help flag to this script.
404
+ # We now keep distinct sets of args, for a cleaner separation of concerns.
405
+
406
+ parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
407
+ if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
408
+ # If we pass only one argument to the script and it's the path to a json file,
409
+ # let's parse it to get our arguments.
410
+ model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
411
+ else:
412
+ model_args, data_args, training_args = parser.parse_args_into_dataclasses()
413
+
414
+ if (
415
+ os.path.exists(training_args.output_dir)
416
+ and os.listdir(training_args.output_dir)
417
+ and training_args.do_train
418
+ and not training_args.overwrite_output_dir
419
+ ):
420
+ raise ValueError(
421
+ f"Output directory ({training_args.output_dir}) already exists and is not empty."
422
+ "Use --overwrite_output_dir to overcome."
423
+ )
424
+
425
+ # Make one log on every process with the configuration for debugging.
426
+ logging.basicConfig(
427
+ format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
428
+ datefmt="%m/%d/%Y %H:%M:%S",
429
+ level=logging.INFO,
430
+ )
431
+ # Setup logging, we only want one process per machine to log things on the screen.
432
+ logger.setLevel(logging.INFO if jax.process_index() == 0 else logging.ERROR)
433
+ if jax.process_index() == 0:
434
+ datasets.utils.logging.set_verbosity_warning()
435
+ transformers.utils.logging.set_verbosity_info()
436
+ else:
437
+ datasets.utils.logging.set_verbosity_error()
438
+ transformers.utils.logging.set_verbosity_error()
439
+
440
+ # Set the verbosity to info of the Transformers logger (on main process only):
441
+ logger.info(f"Training/evaluation parameters {training_args}")
442
+
443
+ # Set seed before initializing model.
444
+ set_seed(training_args.seed)
445
+
446
+ # # Handle the repository creation
447
+ # if training_args.push_to_hub:
448
+ # if training_args.hub_model_id is None:
449
+ # repo_name = get_full_repo_name(
450
+ # Path(training_args.output_dir).absolute().name, token=training_args.hub_token
451
+ # )
452
+ # else:
453
+ # repo_name = training_args.hub_model_id
454
+ # repo = Repository(training_args.output_dir, clone_from=repo_name)
455
+
456
+ # Get the datasets: you can either provide your own CSV/JSON/TXT training and evaluation files (see below)
457
+ # or just provide the name of one of the public datasets available on the hub at https://huggingface.co/datasets/
458
+ # (the dataset will be downloaded automatically from the datasets Hub).
459
+ #
460
+ # For CSV/JSON files, this script will use the column called 'text' or the first column if no column called
461
+ # 'text' is found. You can easily tweak this behavior (see below).
462
+ #
463
+ # In distributed training, the load_dataset function guarantees that only one local process can concurrently
464
+ # download the dataset.
465
+ if data_args.dataset_name is not None:
466
+ # Downloading and loading a dataset from the hub.
467
+ dataset = load_dataset(
468
+ data_args.dataset_name, data_args.dataset_config_name, cache_dir=model_args.cache_dir, keep_in_memory=False
469
+ )
470
+
471
+ if "validation" not in dataset.keys():
472
+ dataset["validation"] = load_dataset(
473
+ data_args.dataset_name,
474
+ data_args.dataset_config_name,
475
+ split=f"train[:{data_args.validation_split_percentage}%]",
476
+ cache_dir=model_args.cache_dir,
477
+ )
478
+ dataset["train"] = load_dataset(
479
+ data_args.dataset_name,
480
+ data_args.dataset_config_name,
481
+ split=f"train[{data_args.validation_split_percentage}%:]",
482
+ cache_dir=model_args.cache_dir,
483
+ )
484
+ else:
485
+ data_files = {}
486
+ dataset_args = {}
487
+ if data_args.train_file is not None:
488
+ data_files["train"] = data_args.train_file
489
+ if data_args.validation_file is not None:
490
+ data_files["validation"] = data_args.validation_file
491
+ extension = data_args.train_file.split(".")[-1]
492
+ if extension == "txt":
493
+ extension = "text"
494
+ dataset_args["keep_linebreaks"] = data_args.keep_linebreaks
495
+ dataset = load_dataset(extension, data_files=data_files, cache_dir=model_args.cache_dir, **dataset_args)
496
+
497
+ if "validation" not in dataset.keys():
498
+ dataset["validation"] = load_dataset(
499
+ extension,
500
+ data_files=data_files,
501
+ split=f"train[:{data_args.validation_split_percentage}%]",
502
+ cache_dir=model_args.cache_dir,
503
+ **dataset_args,
504
+ )
505
+ dataset["train"] = load_dataset(
506
+ extension,
507
+ data_files=data_files,
508
+ split=f"train[{data_args.validation_split_percentage}%:]",
509
+ cache_dir=model_args.cache_dir,
510
+ **dataset_args,
511
+ )
512
+ # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at
513
+ # https://huggingface.co/docs/datasets/loading_datasets.html.
514
+
515
+ # Load pretrained model and tokenizer
516
+
517
+ # Distributed training:
518
+ # The .from_pretrained methods guarantee that only one local process can concurrently
519
+ # download model & vocab.
520
+ if model_args.config_name:
521
+ config = AutoConfig.from_pretrained(model_args.config_name, cache_dir=model_args.cache_dir)
522
+ elif model_args.model_name_or_path:
523
+ config = AutoConfig.from_pretrained(model_args.model_name_or_path, cache_dir=model_args.cache_dir)
524
+ else:
525
+ config = CONFIG_MAPPING[model_args.model_type]()
526
+ logger.warning("You are instantiating a new config instance from scratch.")
527
+
528
+ if model_args.tokenizer_name:
529
+ tokenizer = AutoTokenizer.from_pretrained(
530
+ model_args.tokenizer_name, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_tokenizer
531
+ )
532
+ elif model_args.model_name_or_path:
533
+ tokenizer = AutoTokenizer.from_pretrained(
534
+ model_args.model_name_or_path, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_tokenizer
535
+ )
536
+ else:
537
+ raise ValueError(
538
+ "You are instantiating a new tokenizer from scratch. This is not supported by this script."
539
+ "You can do it from another script, save it, and load it from here, using --tokenizer_name."
540
+ )
541
+
542
+ if model_args.model_name_or_path:
543
+ model = FlaxAutoModelForCausalLM.from_pretrained(
544
+ model_args.model_name_or_path, config=config, seed=training_args.seed, dtype=getattr(jnp, model_args.dtype)
545
+ )
546
+ else:
547
+ model = FlaxAutoModelForCausalLM.from_config(
548
+ config, seed=training_args.seed, dtype=getattr(jnp, model_args.dtype)
549
+ )
550
+
551
+ # Preprocessing the datasets.
552
+ # First we tokenize all the texts.
553
+ if training_args.do_train:
554
+ column_names = dataset["train"].column_names
555
+ else:
556
+ column_names = dataset["validation"].column_names
557
+ text_column_name = "text" if "text" in column_names else column_names[0]
558
+
559
+ # since this will be pickled to avoid _LazyModule error in Hasher force logger loading before tokenize_function
560
+ tok_logger = transformers.utils.logging.get_logger("transformers.tokenization_utils_base")
561
+
562
+ def tokenize_function(examples):
563
+ with CaptureLogger(tok_logger) as cl:
564
+ output = tokenizer(examples[text_column_name])
565
+ # clm input could be much much longer than block_size
566
+ if "Token indices sequence length is longer than the" in cl.out:
567
+ tok_logger.warning(
568
+ "^^^^^^^^^^^^^^^^ Please ignore the warning above - this long input will be chunked into smaller bits before being passed to the model."
569
+ )
570
+ return output
571
+
572
+ tokenized_datasets = dataset.map(
573
+ tokenize_function,
574
+ batched=True,
575
+ num_proc=data_args.preprocessing_num_workers,
576
+ remove_columns=column_names,
577
+ load_from_cache_file=not data_args.overwrite_cache,
578
+ )
579
+
580
+ if data_args.block_size is None:
581
+ block_size = tokenizer.model_max_length
582
+ if block_size > config.max_position_embeddings:
583
+ logger.warning(
584
+ f"The tokenizer picked seems to have a very large `model_max_length` ({tokenizer.model_max_length}). "
585
+ "Picking 1024 instead. You can change that default value by passing --block_size xxx."
586
+ )
587
+ block_size = 1024
588
+ else:
589
+ if data_args.block_size > tokenizer.model_max_length:
590
+ logger.warning(
591
+ f"The block_size passed ({data_args.block_size}) is larger than the maximum length for the model"
592
+ f"({tokenizer.model_max_length}). Using block_size={tokenizer.model_max_length}."
593
+ )
594
+ block_size = min(data_args.block_size, tokenizer.model_max_length)
595
+
596
+ # Main data processing function that will concatenate all texts from our dataset and generate chunks of block_size.
597
+ def group_texts(examples):
598
+ # Concatenate all texts.
599
+ concatenated_examples = {k: list(chain(*examples[k])) for k in examples.keys()}
600
+ total_length = len(concatenated_examples[list(examples.keys())[0]])
601
+ # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can
602
+ # customize this part to your needs.
603
+ if total_length >= block_size:
604
+ total_length = (total_length // block_size) * block_size
605
+ # Split by chunks of max_len.
606
+ result = {
607
+ k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
608
+ for k, t in concatenated_examples.items()
609
+ }
610
+ result["labels"] = result["input_ids"].copy()
611
+ return result
612
+
613
+ # Note that with `batched=True`, this map processes 1,000 texts together, so group_texts throws away a remainder
614
+ # for each of those groups of 1,000 texts. You can adjust that batch_size here but a higher value might be slower
615
+ # to preprocess.
616
+ #
617
+ # To speed up this part, we use multiprocessing. See the documentation of the map method for more information:
618
+ # https://huggingface.co/docs/datasets/package_reference/main_classes.html#datasets.Dataset.map
619
+
620
+ lm_datasets = tokenized_datasets.map(
621
+ group_texts,
622
+ batched=True,
623
+ num_proc=data_args.preprocessing_num_workers,
624
+ load_from_cache_file=not data_args.overwrite_cache,
625
+ )
626
+
627
+ if training_args.do_train:
628
+ if "train" not in tokenized_datasets:
629
+ raise ValueError("--do_train requires a train dataset")
630
+ train_dataset = lm_datasets["train"]
631
+ if data_args.max_train_samples is not None:
632
+ train_dataset = train_dataset.select(range(data_args.max_train_samples))
633
+
634
+ if training_args.do_eval:
635
+ if "validation" not in tokenized_datasets:
636
+ raise ValueError("--do_eval requires a validation dataset")
637
+ eval_dataset = lm_datasets["validation"]
638
+ if data_args.max_eval_samples is not None:
639
+ eval_dataset = eval_dataset.select(range(data_args.max_eval_samples))
640
+
641
+ # Enable tensorboard only on the master node
642
+ has_tensorboard = is_tensorboard_available()
643
+ if has_tensorboard and jax.process_index() == 0:
644
+ try:
645
+ from flax.metrics.tensorboard import SummaryWriter
646
+
647
+ summary_writer = SummaryWriter(log_dir=Path(training_args.output_dir + "/runs"))
648
+ except ImportError as ie:
649
+ has_tensorboard = False
650
+ logger.warning(
651
+ f"Unable to display metrics through TensorBoard because some package are not installed: {ie}"
652
+ )
653
+ else:
654
+ logger.warning(
655
+ "Unable to display metrics through TensorBoard because the package is not installed: "
656
+ "Please run pip install tensorboard to enable."
657
+ )
658
+
659
+ # Initialize our training
660
+ rng = jax.random.PRNGKey(training_args.seed)
661
+ rng, dropout_rng = jax.random.split(rng)
662
+
663
+ # Store some constant
664
+ num_epochs = int(training_args.num_train_epochs)
665
+ train_batch_size = int(training_args.per_device_train_batch_size) * jax.device_count()
666
+ eval_batch_size = int(training_args.per_device_eval_batch_size) * jax.device_count()
667
+ steps_per_epoch = len(train_dataset) // train_batch_size
668
+ total_train_steps = steps_per_epoch * num_epochs
669
+
670
+ # Create learning rate schedule
671
+ linear_decay_lr_schedule_fn = create_learning_rate_fn(
672
+ len(train_dataset),
673
+ train_batch_size,
674
+ training_args.num_train_epochs,
675
+ training_args.warmup_steps,
676
+ training_args.learning_rate,
677
+ )
678
+
679
+ # We use Optax's "masking" functionality to not apply weight decay
680
+ # to bias and LayerNorm scale parameters. decay_mask_fn returns a
681
+ # mask boolean with the same structure as the parameters.
682
+ # The mask is True for parameters that should be decayed.
683
+ # Note that this mask is specifically adapted for FlaxGPT2.
684
+ # For other models, one should correct the layer norm parameter naming
685
+ # accordingly.
686
+ def decay_mask_fn(params):
687
+ flat_params = traverse_util.flatten_dict(params)
688
+ flat_mask = {
689
+ path: (path[-1] != "bias" and path[-2:] not in [("ln_1", "scale"), ("ln_2", "scale"), ("ln_f", "scale")])
690
+ for path in flat_params
691
+ }
692
+ return traverse_util.unflatten_dict(flat_mask)
693
+
694
+ # create adam optimizer
695
+ if training_args.adafactor:
696
+ # We use the default parameters here to initialize adafactor,
697
+ # For more details about the parameters please check https://github.com/deepmind/optax/blob/ed02befef9bf81cbbf236be3d2b0e032e9ed4a40/optax/_src/alias.py#L74
698
+ optimizer = optax.adafactor(
699
+ learning_rate=linear_decay_lr_schedule_fn,
700
+ )
701
+ else:
702
+ optimizer = optax.adamw(
703
+ learning_rate=linear_decay_lr_schedule_fn,
704
+ b1=training_args.adam_beta1,
705
+ b2=training_args.adam_beta2,
706
+ eps=training_args.adam_epsilon,
707
+ weight_decay=training_args.weight_decay,
708
+ mask=decay_mask_fn,
709
+ )
710
+
711
+ # Setup train state
712
+ state = TrainState.create(apply_fn=model.__call__, params=model.params, tx=optimizer, dropout_rng=dropout_rng)
713
+
714
+ if training_args.resume_from_checkpoint:
715
+ state = restore_model_checkpoint(training_args.resume_from_checkpoint, state)
716
+ resume_step = mb_item(state.step)
717
+ if training_args.adafactor:
718
+ state = fake_update(state)
719
+ else:
720
+ resume_step = 0
721
+
722
+ def loss_fn(logits, labels):
723
+ shift_logits = logits[..., :-1, :]
724
+ shift_labels = labels[..., 1:]
725
+ loss = optax.softmax_cross_entropy(shift_logits, onehot(shift_labels, shift_logits.shape[-1]))
726
+ return loss.mean()
727
+
728
+ # Define gradient update step fn
729
+ def train_step(state, batch):
730
+ dropout_rng, new_dropout_rng = jax.random.split(state.dropout_rng)
731
+
732
+ def compute_loss(params):
733
+ labels = batch.pop("labels")
734
+ logits = state.apply_fn(**batch, params=params, dropout_rng=dropout_rng, train=True)[0]
735
+ loss = loss_fn(logits, labels)
736
+ return loss
737
+
738
+ grad_fn = jax.value_and_grad(compute_loss)
739
+ loss, grad = grad_fn(state.params)
740
+ grad = jax.lax.pmean(grad, "batch")
741
+
742
+ new_state = state.apply_gradients(grads=grad, dropout_rng=new_dropout_rng)
743
+
744
+ metrics = {"loss": loss, "learning_rate": linear_decay_lr_schedule_fn(state.step)}
745
+ metrics = jax.lax.pmean(metrics, axis_name="batch")
746
+
747
+ return new_state, metrics
748
+
749
+ # Define eval fn
750
+ def eval_step(params, batch):
751
+ labels = batch.pop("labels")
752
+ logits = model(**batch, params=params, train=False)[0]
753
+ loss = loss_fn(logits, labels)
754
+
755
+ # summarize metrics
756
+ metrics = {"loss": loss}
757
+ metrics = jax.lax.pmean(metrics, axis_name="batch")
758
+ return metrics
759
+
760
+ # Create parallel version of the train and eval step
761
+ p_train_step = jax.pmap(train_step, "batch", donate_argnums=(0,))
762
+ p_eval_step = jax.pmap(eval_step, "batch")
763
+
764
+ # Replicate the train state on each device
765
+ state = state.replicate()
766
+
767
+ logger.info("***** Running training *****")
768
+ logger.info(f" Num examples = {len(train_dataset)}")
769
+ logger.info(f" Num Epochs = {num_epochs}")
770
+ logger.info(f" Num tokenized group examples {len(tokenized_datasets['train'])}")
771
+ logger.info(f" Instantaneous batch size per device = {training_args.per_device_train_batch_size}")
772
+ logger.info(f" Total train batch size (w. parallel & distributed) = {train_batch_size}")
773
+ logger.info(f" Total optimization steps = {total_train_steps}")
774
+
775
+ train_time = 0
776
+ train_metrics = []
777
+ resume_epoch = resume_step // (steps_per_epoch)
778
+ epochs = tqdm(range(num_epochs), desc=f"Epoch ... ({resume_epoch + 1}/{num_epochs})", position=0)
779
+ if resume_step != 0:
780
+ logger.info(f"Skipping to epoch {resume_epoch} step {resume_step}")
781
+ for epoch in epochs:
782
+ # ======================== Training ================================
783
+ if epoch < resume_epoch:
784
+ continue
785
+
786
+ train_start = time.time()
787
+
788
+ # Create sampling rng
789
+ rng, input_rng = jax.random.split(rng)
790
+
791
+ # Generate an epoch by shuffling sampling indices from the train dataset
792
+ train_loader = data_loader(input_rng, train_dataset, train_batch_size, shuffle=True)
793
+ steps_per_epoch = len(train_dataset) // train_batch_size
794
+ # train
795
+ for step in tqdm(range(steps_per_epoch), desc="Training...", position=1, leave=False):
796
+ cur_step = epoch * (len(train_dataset) // train_batch_size) + step
797
+ # skip to the step from which we are resuming
798
+ if cur_step < resume_step:
799
+ continue
800
+
801
+ batch = next(train_loader)
802
+ batch = shard(batch)
803
+ state, train_metric = p_train_step(state, batch)
804
+ train_metrics.append(train_metric)
805
+
806
+
807
+ if cur_step % training_args.logging_steps == 0 and cur_step > 0:
808
+ # Save metrics
809
+ train_metric = unreplicate(train_metric)
810
+ train_time += time.time() - train_start
811
+ if has_tensorboard and jax.process_index() == 0:
812
+ write_train_metric(summary_writer, train_metrics, train_time, cur_step)
813
+
814
+ epochs.write(
815
+ f"Step... ({cur_step} | Loss: {train_metric['loss'].mean()}, Learning Rate: {train_metric['learning_rate'].mean()})"
816
+ )
817
+
818
+ train_metrics = []
819
+
820
+ if cur_step % training_args.eval_steps == 0 and cur_step > 0:
821
+ # ======================== Evaluating ==============================
822
+ eval_metrics = []
823
+ eval_loader = data_loader(input_rng, eval_dataset, eval_batch_size)
824
+ eval_steps = len(eval_dataset) // eval_batch_size
825
+ for _ in tqdm(range(eval_steps), desc="Evaluating...", position=2, leave=False):
826
+ # Model forward
827
+ batch = next(eval_loader)
828
+ batch = shard(batch)
829
+ metrics = p_eval_step(state.params, batch)
830
+ eval_metrics.append(metrics)
831
+
832
+ # normalize eval metrics
833
+ eval_metrics = get_metrics(eval_metrics)
834
+ eval_metrics = jax.tree_map(jnp.mean, eval_metrics)
835
+
836
+ try:
837
+ eval_metrics["perplexity"] = math.exp(eval_metrics["loss"])
838
+ except OverflowError:
839
+ eval_metrics["perplexity"] = float("inf")
840
+
841
+ # Print metrics and update progress bar
842
+ desc = f"Step... ({cur_step} | Eval Loss: {eval_metrics['loss']} | Eval Perplexity: {eval_metrics['perplexity']})"
843
+ epochs.write(desc)
844
+ epochs.desc = desc
845
+
846
+ # Save metrics
847
+ if has_tensorboard and jax.process_index() == 0:
848
+ write_eval_metric(summary_writer, eval_metrics, cur_step)
849
+
850
+ if cur_step % training_args.save_steps == 0 and cur_step > 0:
851
+ # save checkpoint after each epoch and push checkpoint to the hub
852
+ if jax.process_index() == 0:
853
+ save_model_checkpoint(model, training_args.output_dir, state, with_opt=True,
854
+ push_to_hub=training_args.push_to_hub)
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 step {cur_step}", blocking=False)
860
+
861
+ # Eval after training
862
+ if training_args.do_eval:
863
+ eval_metrics = []
864
+ eval_loader = data_loader(input_rng, eval_dataset, eval_batch_size)
865
+ eval_steps = len(eval_dataset) // eval_batch_size
866
+ for _ in tqdm(range(eval_steps), desc="Evaluating...", position=2, leave=False):
867
+ # Model forward
868
+ batch = shard(next(eval_loader))
869
+ metrics = p_eval_step(state.params, batch)
870
+ eval_metrics.append(metrics)
871
+
872
+ # normalize eval metrics
873
+ eval_metrics = get_metrics(eval_metrics)
874
+ eval_metrics = jax.tree_map(lambda x: jnp.mean(x).item(), eval_metrics)
875
+
876
+ try:
877
+ eval_metrics["perplexity"] = math.exp(eval_metrics["loss"])
878
+ except OverflowError:
879
+ eval_metrics["perplexity"] = float("inf")
880
+
881
+ if jax.process_index() == 0:
882
+ eval_metrics = {f"eval_{metric_name}": value for metric_name, value in eval_metrics.items()}
883
+ path = os.path.join(training_args.output_dir, "eval_results.json")
884
+ with open(path, "w") as f:
885
+ json.dump(eval_metrics, f, indent=4, sort_keys=True)
886
+
887
+ # save model after training is over
888
+ if jax.process_index() == 0:
889
+ save_model_checkpoint(model, training_args.output_dir, state, with_opt=False,
890
+ push_to_hub=training_args.push_to_hub)
891
+
892
+
893
+
894
+ if __name__ == "__main__":
895
+ main()
run_gpt.sh ADDED
@@ -0,0 +1,37 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/bin/bash
2
+
3
+ export HF_PROJECT="gpt2-large-dutch-2"
4
+
5
+ # Variables for training the tokenizer and creating the config
6
+ export VOCAB_SIZE="50257"
7
+ export DATASET="yhavinga/mc4_nl_cleaned" # Name of the dataset in the Huggingface Hub
8
+ export DATASET_CONFIG="full" # Config of the dataset in the Huggingface Hub
9
+ export DATASET_SPLIT="train" # Split to use for training tokenizer and model
10
+ export TEXT_FIELD="text" # Field containing the text to be used for training
11
+ export CONFIG_TYPE="gpt2-large" # Config that our model will use
12
+ export MODEL_PATH="${HOME}/data/${HF_PROJECT}" # Path to the model, e.g. here inside the mount
13
+
14
+ python run_clm_flax.py \
15
+ --output_dir="${MODEL_PATH}" \
16
+ --model_type="gpt2" \
17
+ --config_name="${MODEL_PATH}" \
18
+ --model_name_or_path="${MODEL_PATH}" \
19
+ --tokenizer_name="${MODEL_PATH}" \
20
+ --preprocessing_num_workers="96" \
21
+ --do_train --do_eval \
22
+ --dataset_name="${DATASET}" \
23
+ --dataset_config_name="${DATASET_CONFIG}" \
24
+ --block_size="512" \
25
+ --per_device_train_batch_size="4" \
26
+ --per_device_eval_batch_size="4" \
27
+ --learning_rate="0.000033" --warmup_steps="5000" \
28
+ --adafactor \
29
+ --overwrite_output_dir \
30
+ --num_train_epochs="1" \
31
+ --logging_steps="500" \
32
+ --save_steps="20000" \
33
+ --eval_steps="2500"
34
+
35
+ # \
36
+ # --push_to_hub
37
+ # --adam_beta1="0.9" --adam_beta2="0.98" --weight_decay="0.01" \
runs/events.out.tfevents.1641027824.t1v-n-f9cfcc28-w-0.77645.0.v2 DELETED
@@ -1,3 +0,0 @@
1
- version https://git-lfs.github.com/spec/v1
2
- oid sha256:d605b44ebcfa0fd3d225ad471f2e074ac95dca362f4e2b8c21d0f64e634ab2ce
3
- size 40
 
 
 
 
runs/events.out.tfevents.1641054222.t1v-n-f9cfcc28-w-0.98674.0.v2 DELETED
@@ -1,3 +0,0 @@
1
- version https://git-lfs.github.com/spec/v1
2
- oid sha256:231cba331f1db1eca87731face3cc13ff79c93c16454d3eb61a401444ebc65ab
3
- size 40
 
 
 
 
runs/events.out.tfevents.1641055087.t1v-n-f9cfcc28-w-0.110334.0.v2 DELETED
@@ -1,3 +0,0 @@
1
- version https://git-lfs.github.com/spec/v1
2
- oid sha256:acfd659bc685866d4a47dc94a305ef2138f16b9abdbe1bafb46903f2b24c4b69
3
- size 40
 
 
 
 
runs/events.out.tfevents.1641055391.t1v-n-f9cfcc28-w-0.112189.0.v2 CHANGED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:1f49cb4b742ac7ebe1f71281334586064571b27c97412cd1d4ddbae2852e08da
3
- size 53678253
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:a583ef507312253e62d4ea82ed51917a55ccf61c10a4e41ac592dd47cf94fcc5
3
+ size 53752823
special_tokens_map.json ADDED
@@ -0,0 +1 @@
 
 
1
+ {"bos_token": "<|endoftext|>", "eos_token": "<|endoftext|>", "unk_token": "<|endoftext|>"}
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"}
vocab.json ADDED
The diff for this file is too large to render. See raw diff