Merge pull request #2 from tlkh/finetuning

Added instructions and script for distributed training with Horovod

README.md

@@ -97,14 +97,26 @@ python3 src/interactive_conditional_samples.py -- --help
 To retrain GPT-2 117M model on a custom text dataset:
 
 ```
-PYTHONPATH=src ./train --dataset <file|directory|glob>
+PYTHONPATH=src ./train.py --dataset <file|directory|glob>
 ```
 
 If you want to precompute the dataset's encoding for multiple runs, you can instead use:
 
 ```
 PYTHONPATH=src ./encode.py <file|directory|glob> /path/to/encoded.npz
-PYTHONPATH=src ./train --dataset /path/to/encoded.npz
+PYTHONPATH=src ./train.py --dataset /path/to/encoded.npz
+```
+
+To do distributed on multiple GPUs or machines using Horovod: 
+
+```
+mpirun -np 4 \
+ -H localhost:4 \
+ -bind-to none -map-by slot \
+ -x NCCL_DEBUG=INFO -x LD_LIBRARY_PATH -x PATH \
+ -x PYTHONPATH=src \
+ -mca pml ob1 -mca btl ^openib \
+ /home/jovyan/gpt-2/train-horovod.py --dataset encoded.npz
 ```
 
 ## GPT-2 samples

train-horovod.py

@@ -0,0 +1,260 @@
+#!/usr/bin/env python3
+# Usage:
+# PYTHONPATH=src ./train --dataset <file|directory|glob>
+
+import fire
+import json
+import os
+import numpy as np
+import tensorflow as tf
+import random
+import time
+
+import horovod.tensorflow as hvd
+
+import model, sample, encoder
+
+CHECKPOINT_DIR = 'checkpoint'
+SAMPLE_DIR = 'samples'
+
+hvd.init()
+
+def maketree(path):
+ try:
+ os.makedirs(path)
+ except:
+ pass
+
+
+def load_dataset(enc, path):
+ paths = []
+ if os.path.isfile(path):
+ # Simple file
+ paths.append(path)
+ elif os.path.isdir(path):
+ # Directory
+ for (dirpath, _, fnames) in os.walk(path):
+ for fname in fnames:
+ paths.append(os.path.join(dirpath, fname))
+ else:
+ # Assume glob
+ paths = glob.glob(path)
+
+ token_chunks = []
+ for path in paths:
+ print(str(hvd.local_rank()), 'Reading', path)
+ if path.endswith('.npz'):
+ # Pre-encoded
+ with np.load(path) as npz:
+ for item in npz.files:
+ token_chunks.append(npz[item])
+ else:
+ with open(path, 'r') as fp:
+ raw_text = fp.read()
+ tokens = np.stack(enc.encode(raw_text))
+ token_chunks.append(tokens)
+ return token_chunks
+
+
+def binary_search(f, lo, hi):
+ if f(lo) or not f(hi):
+ return None
+ while hi > lo + 1:
+ mid = (lo + hi) // 2
+ if f(mid):
+ hi = mid
+ else:
+ lo = mid
+ return hi
+
+
+class Sampler(object):
+ """Fairly samples a slice from a set of variable sized chunks.
+
+ 'Fairly' means that the distribution is the same as sampling from one concatenated chunk,
+ but without crossing chunk boundaries."""
+
+ def __init__(self, chunks):
+ self.chunks = chunks
+ self.total_size = sum(chunk.shape[0] for chunk in chunks)
+ self.boundaries = [0]
+ for i in range(len(chunks)):
+ self.boundaries.append(self.boundaries[-1] + chunks[i].shape[0])
+
+ def sample(self, length):
+ assert length < self.total_size // len(
+ self.chunks
+ ), "Dataset files are too small to sample {} tokens at a time".format(length)
+ while True:
+ index = random.randint(0, self.total_size - length - 1)
+ i = binary_search(lambda j: self.boundaries[j] > index, 0,
+ len(self.boundaries) - 1) - 1
+ if self.boundaries[i + 1] > index + length:
+ within_chunk = index - self.boundaries[i]
+ return self.chunks[i][within_chunk:within_chunk + length]
+
+
+def train_main(dataset,
+ model_name='117M',
+ seed=None,
+ batch_size=2,
+ sample_length=1023,
+ sample_num=1,
+ sample_every=4500,
+ run_name='run1',
+ restore_from='latest',
+ save_every=2000):
+
+ enc = encoder.get_encoder(model_name)
+ hparams = model.default_hparams()
+ with open(os.path.join('models', model_name, 'hparams.json')) as f:
+ hparams.override_from_dict(json.load(f))
+
+ if sample_length is None:
+ sample_length = hparams.n_ctx // 2
+ elif sample_length > hparams.n_ctx:
+ raise ValueError(
+ "Can't get samples longer than window size: %s" % hparams.n_ctx)
+
+ # TF config
+
+ config = tf.ConfigProto()
+ config.gpu_options.visible_device_list = str(hvd.local_rank())
+ config.gpu_options.allow_growth = True
+
+ with tf.Session(config=config) as sess:
+ context = tf.placeholder(tf.int32, [batch_size, None])
+ np.random.seed(seed)
+ tf.set_random_seed(seed)
+ output = model.model(hparams=hparams, X=context)
+ loss = tf.reduce_mean(
+ tf.nn.sparse_softmax_cross_entropy_with_logits(
+ labels=context[:, 1:], logits=output['logits'][:, :-1]))
+
+ tf_sample = sample.sample_sequence(
+ hparams=hparams,
+ length=sample_length,
+ context=context,
+ batch_size=batch_size,
+ temperature=0.8,
+ top_k=40)
+
+ train_vars = [v for v in tf.trainable_variables() if 'model' in v.name]
+
+ opt = tf.train.AdamOptimizer()
+ opt = hvd.DistributedOptimizer(opt)
+ train_op = opt.minimize(loss, var_list=train_vars)
+
+ # Horovod: broadcast initial variable states from rank 0 to all other processes.
+ # This is necessary to ensure consistent initialization of all workers when
+ # training is started with random weights or restored from a checkpoint.
+ bcast = hvd.broadcast_global_variables(0)
+
+ saver = tf.train.Saver(
+ var_list=train_vars,
+ max_to_keep=5,
+ keep_checkpoint_every_n_hours=2)
+
+ sess.run(tf.global_variables_initializer())
+
+
+ if restore_from == 'latest':
+ ckpt = tf.train.latest_checkpoint(
+ os.path.join(CHECKPOINT_DIR, run_name))
+ if ckpt is None:
+ # Get fresh GPT weights if new run.
+ ckpt = tf.train.latest_checkpoint(
+ os.path.join('models', model_name))
+ elif restore_from == 'fresh':
+ ckpt = tf.train.latest_checkpoint(
+ os.path.join('models', model_name))
+ else:
+ ckpt = tf.train.latest_checkpoint(restore_from)
+ print(str(hvd.local_rank()), 'Loading checkpoint', ckpt)
+ saver.restore(sess, ckpt)
+
+ bcast.run()
+
+ print(str(hvd.local_rank()), 'Loading dataset...')
+ chunks = load_dataset(enc, dataset)
+ data_sampler = Sampler(chunks)
+ print(str(hvd.local_rank()), 'dataset has', data_sampler.total_size, 'tokens')
+ print(str(hvd.local_rank()), 'Training...')
+
+ counter = 1
+ if os.path.exists(os.path.join(CHECKPOINT_DIR, run_name, 'counter')):
+ # Load the step number if we're resuming a run
+ # Add 1 so we don't immediately try to save again
+ with open(os.path.join(CHECKPOINT_DIR, run_name, 'counter'),
+ 'r') as fp:
+ counter = int(fp.read()) + 1
+
+ def save():
+ maketree(os.path.join(CHECKPOINT_DIR, run_name))
+ print(
+ 'Saving',
+ os.path.join(CHECKPOINT_DIR, run_name,
+ 'model-{}').format(counter))
+ saver.save(
+ sess,
+ os.path.join(CHECKPOINT_DIR, run_name, 'model'),
+ global_step=counter)
+ with open(os.path.join(CHECKPOINT_DIR, run_name, 'counter'),
+ 'w') as fp:
+ fp.write(str(counter) + '\n')
+
+ def generate_samples():
+ context_tokens = data_sampler.sample(1)
+ all_text = []
+ index = 0
+ while index < sample_num:
+ out = sess.run(
+ tf_sample, feed_dict={context: batch_size*[context_tokens]})
+ for i in range(min(sample_num - index, batch_size)):
+ text = enc.decode(out[i])
+ text = '======== SAMPLE {} ========\n{}\n'.format(index + 1, text)
+ all_text.append(text)
+ index += 1
+ print(text)
+ maketree(os.path.join(SAMPLE_DIR, run_name))
+ with open(
+ os.path.join(SAMPLE_DIR, run_name,
+ 'samples-{}').format(counter), 'w') as fp:
+ fp.write('\n'.join(all_text))
+
+ avg_loss = (0.0, 0.0)
+ start_time = time.time()
+
+ try:
+ while True:
+
+ batch = [data_sampler.sample(1024) for _ in range(batch_size)]
+
+ _, lv = sess.run((train_op, loss), feed_dict={context: batch})
+
+ avg_loss = (avg_loss[0] * 0.99 + lv, avg_loss[1] * 0.99 + 1.0)
+
+ if hvd.rank() == 0:
+ if counter % save_every == 0:
+ save()
+ if counter % sample_every == 0:
+ generate_samples()
+
+ print(
+ '[{counter} | {time:2.2f}] loss={loss:2.2f} avg={avg:2.2f}'
+ .format(
+ counter=counter,
+ time=time.time() - start_time,
+ loss=lv,
+ avg=avg_loss[0] / avg_loss[1]))
+
+ counter += 1
+
+ except KeyboardInterrupt:
+ print('interrupted')
+ if hvd.rank() == 0:
+ save()
+
+
+if __name__ == '__main__':
+ fire.Fire(train_main)