chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,12 @@
|
||||
"""
|
||||
---
|
||||
title: Samples
|
||||
summary: >
|
||||
Samples for inference and fine-tuning
|
||||
---
|
||||
|
||||
# Samples
|
||||
|
||||
* [Generating text](generate.html)
|
||||
* [Fine tuning the biases with pipeline-parallel training](finetune.html)
|
||||
"""
|
||||
@@ -0,0 +1,128 @@
|
||||
"""
|
||||
---
|
||||
title: Fine Tune GPT-NeoX
|
||||
summary: >
|
||||
Fine tune GPT-NeoX biases with Fairscale pipeline parallel module
|
||||
---
|
||||
|
||||
# Fine Tune GPT-NeoX
|
||||
|
||||
This shows how to fine tune GPT-NeoX with pipeline parallelism.
|
||||
"""
|
||||
|
||||
import fairscale
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.utils.data
|
||||
import torch.utils.data
|
||||
import typing
|
||||
from torch.utils.data import DataLoader, RandomSampler
|
||||
|
||||
from labml import experiment, monit, tracker, lab
|
||||
from labml.configs import option
|
||||
from labml.logger import inspect
|
||||
from labml_nn.neox.utils.text_dataset import get_training_data
|
||||
from labml_nn.neox.utils.finetune import FineTuneBiases
|
||||
from labml_nn.neox.model import LayerGenerator, NeoXModule
|
||||
from labml_nn.neox.utils import balance_layers_simple
|
||||
from labml_nn.neox.utils.trainer import PipelineParallelTrainerConf
|
||||
|
||||
|
||||
@option(PipelineParallelTrainerConf.layers, 'PipelineBiases')
|
||||
def neox_layers(c: PipelineParallelTrainerConf):
|
||||
"""
|
||||
### Load GPT-NeoX layers
|
||||
"""
|
||||
return list(LayerGenerator(is_clone_layers=c.is_clone_layers,
|
||||
filter_layers=c.filter_layers,
|
||||
dtype=c.dtype,
|
||||
).load())
|
||||
|
||||
|
||||
@option(PipelineParallelTrainerConf.fine_tuner, 'PipelineBiases')
|
||||
def fine_tune_biases(c: PipelineParallelTrainerConf):
|
||||
"""
|
||||
### Create fine tuner for biases
|
||||
"""
|
||||
|
||||
fine_tuner = FineTuneBiases(typing.cast(typing.List[NeoXModule], c.layers))
|
||||
# Mark biases as trainable
|
||||
fine_tuner.set_trainable_params()
|
||||
|
||||
#
|
||||
return fine_tuner
|
||||
|
||||
|
||||
@option(PipelineParallelTrainerConf.model, 'PipelineBiases')
|
||||
def pipe_model(c: PipelineParallelTrainerConf):
|
||||
"""
|
||||
### Create pipeline parallel model
|
||||
"""
|
||||
|
||||
if c.is_checkpointing:
|
||||
raise NotImplementedError()
|
||||
else:
|
||||
layers = c.layers
|
||||
|
||||
# Make sure the finetuner is initialized
|
||||
_ = c.fine_tuner
|
||||
|
||||
# Create the Pipe module
|
||||
with monit.section('Pipe'):
|
||||
# Get the layer distribution across GPUs
|
||||
balance = balance_layers_simple(len(layers), c.n_gpus)
|
||||
inspect(balance=balance)
|
||||
# Devices for each GPU
|
||||
devices = [torch.device(f'cuda:{i}') for i in range(c.n_gpus)]
|
||||
# Create Fairscale Pipe module
|
||||
pipe_model = fairscale.nn.Pipe(nn.Sequential(*layers),
|
||||
balance=balance,
|
||||
devices=devices,
|
||||
chunks=c.chunks)
|
||||
|
||||
#
|
||||
return pipe_model
|
||||
|
||||
|
||||
@option(PipelineParallelTrainerConf.train_loader)
|
||||
def tiny_shakespeare(c: PipelineParallelTrainerConf):
|
||||
"""
|
||||
#### Tiny Shakespeare dataset
|
||||
"""
|
||||
dataset = get_training_data(c.max_seq_len)
|
||||
|
||||
return DataLoader(dataset,
|
||||
batch_size=c.batch_size,
|
||||
sampler=RandomSampler(dataset, replacement=True))
|
||||
|
||||
|
||||
def main():
|
||||
# Create experiment
|
||||
experiment.create(name='pipe_neox_biases',
|
||||
writers={'screen', 'web_api'})
|
||||
|
||||
# Initialize configs
|
||||
conf = PipelineParallelTrainerConf()
|
||||
experiment.configs(conf, {
|
||||
'learning_rate': 3e-4,
|
||||
'is_checkpointing': False,
|
||||
'max_seq_len': 128,
|
||||
'batch_size': 64,
|
||||
'chunks': 8,
|
||||
})
|
||||
|
||||
# Start the experiment
|
||||
with experiment.start():
|
||||
# Initialize the model. Do this before the loop for cleaner logs.
|
||||
_ = conf.model
|
||||
|
||||
# Train
|
||||
for epoch in monit.loop(conf.epochs):
|
||||
conf.train_epoch()
|
||||
tracker.new_line()
|
||||
torch.save(conf.fine_tuner.state_dict(), str(lab.get_data_path() / 'fine_tune.pt'))
|
||||
|
||||
|
||||
#
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,102 @@
|
||||
"""
|
||||
---
|
||||
title: Generate Text with GPT-NeoX
|
||||
summary: >
|
||||
Generate Text with GPT-NeoX
|
||||
---
|
||||
|
||||
# Generate Text with GPT-NeoX
|
||||
|
||||
This shows how to generate text from GPT-NeoX with a single GPU.
|
||||
|
||||
This needs a GPU with more than 45GB memory.
|
||||
"""
|
||||
|
||||
# Imports
|
||||
from typing import List
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from labml import monit
|
||||
from labml_nn.neox.model import LayerGenerator
|
||||
from labml_nn.neox.utils import get_tokens, print_tokens
|
||||
from labml_nn.neox.utils.cache import get_cache
|
||||
|
||||
# List of layers to load. This is used for testing.
|
||||
# You can assign a subset of layers like `{0, 1}` so that it only loads
|
||||
# the first to transformer layers.
|
||||
LAYERS = None
|
||||
|
||||
# Prompt to complete
|
||||
PROMPT = 'Einstein was born in the German Empire, but moved to Switzerland in 1895, forsaking his German'
|
||||
|
||||
|
||||
def infer(model: nn.Module, ids: List[int], device: torch.device):
|
||||
"""
|
||||
### Predict the next token
|
||||
|
||||
:param model: is the model
|
||||
:param ids: are the input token ids
|
||||
:param device: is the device of the model
|
||||
"""
|
||||
|
||||
with torch.no_grad():
|
||||
# Get the tokens
|
||||
x = torch.tensor(ids)[None, :].to(device)
|
||||
# Eval model
|
||||
x = model(x)
|
||||
|
||||
# Return predicted token
|
||||
return x[0].max(dim=-1)[1].tolist()
|
||||
|
||||
|
||||
def generate():
|
||||
"""
|
||||
## Generate text
|
||||
"""
|
||||
|
||||
# Setup [cache](../utils/cache.html) to cache intermediate key/value pairs for faster generation
|
||||
cache = get_cache()
|
||||
cache.set('use_cache', True)
|
||||
|
||||
# Device
|
||||
device = torch.device('cuda:0')
|
||||
|
||||
# Load layers
|
||||
layers = list(LayerGenerator(is_clone_layers=True,
|
||||
filter_layers=LAYERS,
|
||||
dtype=torch.float16,
|
||||
device=device,
|
||||
).load())
|
||||
|
||||
model = nn.Sequential(*layers)
|
||||
|
||||
# Get token ids
|
||||
ids = get_tokens(PROMPT)
|
||||
|
||||
# Run the model
|
||||
cache.set('state_ids', (None, 1))
|
||||
with monit.section('Infer'):
|
||||
next_token = infer(model, ids, device)[-1]
|
||||
|
||||
# Append the predicted token
|
||||
ids += [next_token]
|
||||
|
||||
# Predict 100 tokens
|
||||
for i in range(1, 100):
|
||||
# Set the state to use cached activations
|
||||
cache.set('state_ids', (i, i + 1))
|
||||
# Get next token. Note that we only feed the last token to the model because
|
||||
# we cache the key/value pairs of previous tokens.
|
||||
with monit.section('Infer'):
|
||||
next_token = infer(model, [next_token], device)[-1]
|
||||
# Append the predicted token
|
||||
ids += [next_token]
|
||||
# Print
|
||||
print_tokens(ids, [ids])
|
||||
|
||||
|
||||
#
|
||||
if __name__ == '__main__':
|
||||
generate()
|
||||
@@ -0,0 +1,91 @@
|
||||
"""
|
||||
---
|
||||
title: Generate Text with GPT-NeoX using LLM.int8() quantization
|
||||
summary: >
|
||||
Generate Text with GPT-NeoX using LLM.int8() quantization
|
||||
---
|
||||
|
||||
# Generate Text with GPT-NeoX using LLM.int8() quantization
|
||||
|
||||
This shows how to generate text from GPT-NeoX using [LLM.int8() quantization](../utils/llm_int8.html).
|
||||
|
||||
This needs a GPU with 24GB memory.
|
||||
"""
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from labml import monit
|
||||
from labml_nn.neox.model import LayerGenerator
|
||||
from labml_nn.neox.samples.generate import PROMPT, infer
|
||||
from labml_nn.neox.utils import get_tokens, print_tokens
|
||||
from labml_nn.neox.utils.cache import get_cache
|
||||
|
||||
|
||||
def generate():
|
||||
"""
|
||||
## Generate text
|
||||
"""
|
||||
|
||||
# Setup [cache](../utils/cache.html) to cache intermediate key/value pairs for faster generation
|
||||
cache = get_cache()
|
||||
cache.set('use_cache', True)
|
||||
|
||||
# Device
|
||||
device = torch.device('cuda:0')
|
||||
|
||||
# Load layers in float16 into CPU. We convert the layers to int8 later, because doing that
|
||||
# on the fly after loading layers to GPU causes CUDA memory fragmentation
|
||||
# (about 3GB memory can get lost due to fragmentation).
|
||||
layer_generator = LayerGenerator(is_clone_layers=True,
|
||||
dtype=torch.float16,
|
||||
device=torch.device('cpu'),
|
||||
is_llm_int8=False,
|
||||
)
|
||||
layers = list(layer_generator.load())
|
||||
|
||||
# This reduces CUDA memory fragmentation
|
||||
for layer in monit.iterate('Convert to int8', layers, is_children_silent=True):
|
||||
layer_generator.post_load_prepare(layer,
|
||||
device=device,
|
||||
is_llm_int8=True,
|
||||
llm_int8_threshold=6.0,
|
||||
)
|
||||
layer.to(device)
|
||||
|
||||
# Create `nn.Sequential` model
|
||||
model = nn.Sequential(*layers)
|
||||
|
||||
# Clear cache and print memory summary for debugging
|
||||
torch.cuda.empty_cache()
|
||||
print(torch.cuda.memory_summary())
|
||||
|
||||
# Get token ids
|
||||
ids = get_tokens(PROMPT)
|
||||
|
||||
# Run the model.
|
||||
# We use the [`infer`](generate.html) function defined in [`generate.py`](generate.html)
|
||||
cache.set('state_ids', (None, 1))
|
||||
with monit.section('Infer'):
|
||||
next_token = infer(model, ids, device)[-1]
|
||||
|
||||
# Append the predicted token
|
||||
ids += [next_token]
|
||||
|
||||
# Predict 100 tokens
|
||||
for i in range(1, 100):
|
||||
# Set the state to use cached activations
|
||||
cache.set('state_ids', (i, i + 1))
|
||||
# Get next token. Note that we only feed the last token to the model because
|
||||
# we cache the key/value pairs of previous tokens.
|
||||
with monit.section('Infer'):
|
||||
next_token = infer(model, [next_token], device)[-1]
|
||||
# Append the predicted token
|
||||
ids += [next_token]
|
||||
# Print
|
||||
print_tokens(ids, [ids])
|
||||
|
||||
|
||||
#
|
||||
if __name__ == '__main__':
|
||||
generate()
|
||||
Reference in New Issue
Block a user