chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,134 @@
|
||||
"""
|
||||
---
|
||||
title: Utilities and Helpers
|
||||
summary: >
|
||||
Utilities and helper functions
|
||||
---
|
||||
|
||||
# Utilities and Helpers
|
||||
|
||||
* [Cache for intermediate activations (for faster inference)](cache.html)
|
||||
* [Tools for finetuning](finetune.html)
|
||||
* [Trainer](trainer.html)
|
||||
* [Text dataset](text_dataset.html)
|
||||
"""
|
||||
import typing
|
||||
from typing import List, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from labml import logger
|
||||
from labml.logger import Text
|
||||
from labml_nn.neox.tokenizer import get_tokenizer
|
||||
|
||||
if typing.TYPE_CHECKING:
|
||||
from tokenizers import Tokenizer
|
||||
|
||||
# Tokenizer singleton
|
||||
_TOKENIZER: Optional['Tokenizer'] = None
|
||||
|
||||
|
||||
def get_tokens(text: str) -> List[int]:
|
||||
"""
|
||||
### Get token ids
|
||||
|
||||
:param text: is the text to tokenize
|
||||
:return: the token ids
|
||||
"""
|
||||
global _TOKENIZER
|
||||
if _TOKENIZER is None:
|
||||
_TOKENIZER = get_tokenizer()
|
||||
return _TOKENIZER.encode_batch([text])[0].ids
|
||||
|
||||
|
||||
def print_token_outputs(ids: List[int], *xs: torch.Tensor):
|
||||
"""
|
||||
### Print tokens from model outputs
|
||||
|
||||
Pretty prints target tokens along side outputs from the model(s).
|
||||
|
||||
:param ids: are the target token ids
|
||||
:param xs: are the model(s) outputs
|
||||
"""
|
||||
ids = ids + [-1]
|
||||
xs = [[-1] + x[0].max(dim=-1)[1].tolist() for x in xs]
|
||||
|
||||
print_tokens(ids, xs)
|
||||
|
||||
|
||||
def print_tokens(target: List[int], others: List[List[int]]):
|
||||
"""
|
||||
### Print tokens
|
||||
|
||||
Pretty prints tokens for comparison
|
||||
|
||||
:param target: are the target token ids
|
||||
:param others: are the sampled outputs from the model(s)
|
||||
"""
|
||||
|
||||
# Load tokenizer
|
||||
global _TOKENIZER
|
||||
if _TOKENIZER is None:
|
||||
_TOKENIZER = get_tokenizer()
|
||||
|
||||
# Convert the tokens to list of strings
|
||||
text = []
|
||||
for i in range(len(target)):
|
||||
tokens = [_TOKENIZER.decode([target[i]]) if target[i] != -1 else '---']
|
||||
for j in range(len(others)):
|
||||
tokens.append(_TOKENIZER.decode([others[j][i]]) if others[j][i] != -1 else '---')
|
||||
|
||||
text.append(tokens)
|
||||
|
||||
# Stats
|
||||
correct = [0 for _ in others]
|
||||
total = 0
|
||||
|
||||
# Iterate through tokens
|
||||
for i in range(len(target)):
|
||||
parts = [(f'{i}: ', Text.meta)]
|
||||
parts += [('"', Text.subtle), (text[i][0], Text.subtle), ('"', Text.subtle), '\t']
|
||||
|
||||
# Empty target
|
||||
if target[i] == -1:
|
||||
for j in range(len(others)):
|
||||
parts += [('"', Text.subtle), (text[i][j + 1], Text.subtle), ('"', Text.subtle), '\t']
|
||||
|
||||
logger.log(parts)
|
||||
continue
|
||||
|
||||
# Number of tokens
|
||||
total += 1
|
||||
|
||||
# Other outputs
|
||||
for j in range(len(others)):
|
||||
correct[j] += 1 if others[j][i] == target[i] else 0
|
||||
|
||||
parts += [('"', Text.subtle),
|
||||
(text[i][j + 1], Text.success if others[j][i] == target[i] else Text.danger),
|
||||
('"', Text.subtle), '\t']
|
||||
|
||||
logger.log(parts)
|
||||
|
||||
# Stats
|
||||
parts = [(f'{total}', Text.highlight), '\t']
|
||||
for j in range(len(others)):
|
||||
parts += [(f'{correct[j]}', Text.value), '\t']
|
||||
logger.log(parts)
|
||||
|
||||
|
||||
def balance_layers_simple(n_layers: int, n_chunks: int):
|
||||
"""
|
||||
### Balance layers
|
||||
|
||||
Split the `n_layers` into `n_chunks`. This is used for pipeline parallel training.
|
||||
|
||||
:param n_layers: is the number of layers
|
||||
:param n_chunks: is the number of chunks
|
||||
:return: returns a list with the number of layers for each chunk
|
||||
"""
|
||||
balance = []
|
||||
for i in range(n_chunks):
|
||||
balance.append((n_layers - sum(balance)) // (n_chunks - i))
|
||||
|
||||
return list(reversed(balance))
|
||||
@@ -0,0 +1,118 @@
|
||||
"""
|
||||
---
|
||||
title: Cache for Intermediate Activations
|
||||
summary: >
|
||||
Cache for intermediate activations for faster inference.
|
||||
---
|
||||
|
||||
# Cache for Intermediate Activations
|
||||
|
||||
During inference the model outputs token by token.
|
||||
We use this simple cache to store key's and value's attention layers,
|
||||
so that we don't have to recompute them for previous tokens.
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
|
||||
class Cache:
|
||||
"""
|
||||
## Cache
|
||||
|
||||
This maintains a key-value cache and queues push values and pop them in the same order.
|
||||
The queues are useful since we have multiple attention layers.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self._cache = {}
|
||||
|
||||
def clear_all(self):
|
||||
"""
|
||||
### Clear cache
|
||||
"""
|
||||
self._cache = {}
|
||||
|
||||
def push(self, name: str, value: Any):
|
||||
"""
|
||||
### Push a value to a queue
|
||||
|
||||
:param name: is the name of the queue
|
||||
:param value: is the value to be pushed
|
||||
"""
|
||||
|
||||
# Create an empty queue if it's not present
|
||||
if name not in self._cache:
|
||||
self._cache[name] = []
|
||||
|
||||
# Push to the queue
|
||||
self._cache[name].append(value)
|
||||
|
||||
def q_size(self, name):
|
||||
"""
|
||||
### Return the size of the queue
|
||||
|
||||
:param name: is the name of the queue
|
||||
:return: size of the queue if exists else None
|
||||
"""
|
||||
|
||||
if name not in self._cache:
|
||||
return None
|
||||
|
||||
if type(self._cache[name]) != list:
|
||||
return None
|
||||
|
||||
return len(self._cache[name])
|
||||
|
||||
def pop(self, name: str):
|
||||
"""
|
||||
### Pop from a queue
|
||||
|
||||
:param name: is the name of the queue
|
||||
:return: the value
|
||||
"""
|
||||
return self._cache[name].pop(0)
|
||||
|
||||
def set(self, key: str, value: Any):
|
||||
"""
|
||||
### Cache a value
|
||||
|
||||
:param key: is the name of the value to be cached
|
||||
:param value: is the value
|
||||
"""
|
||||
self._cache[key] = value
|
||||
|
||||
def get(self, key: str, default: Any = None):
|
||||
"""
|
||||
### Retrieve a value from cache
|
||||
|
||||
:param key: is the name used when caching
|
||||
:param default: is the default value if the cache is empty
|
||||
:return: the cached value
|
||||
"""
|
||||
return self._cache.get(key, default)
|
||||
|
||||
def clear(self, key: str):
|
||||
"""
|
||||
### Clear a cache value
|
||||
|
||||
:param key: is the name used when caching
|
||||
"""
|
||||
del self._cache[key]
|
||||
|
||||
|
||||
# Singleton for cache
|
||||
_INSTANCE = None
|
||||
|
||||
|
||||
def get_cache() -> Cache:
|
||||
"""
|
||||
### Get the cache instance
|
||||
|
||||
:return: the cache instance
|
||||
"""
|
||||
global _INSTANCE
|
||||
|
||||
if _INSTANCE is None:
|
||||
_INSTANCE = Cache()
|
||||
|
||||
return _INSTANCE
|
||||
@@ -0,0 +1,55 @@
|
||||
from typing import List, Dict
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from labml_nn.neox.model import TransformerLayer, NeoXModule
|
||||
|
||||
|
||||
class FineTuner:
|
||||
def __init__(self, layers: List[NeoXModule]):
|
||||
self.layers = layers
|
||||
|
||||
def get_trainable_params(self) -> Dict[str, nn.Parameter]:
|
||||
params = {}
|
||||
for i, layer in enumerate(self.layers):
|
||||
params.update(self.get_layer_trainable_params(layer, prefix=f'layer_{i :02d}'))
|
||||
|
||||
return params
|
||||
|
||||
def get_layer_trainable_params(self, layer: NeoXModule, prefix: str) -> Dict[str, nn.Parameter]:
|
||||
raise NotImplementedError
|
||||
|
||||
def set_trainable_params(self):
|
||||
for layer in self.layers:
|
||||
# Set `requires_grad` to `False` for the entire layer.
|
||||
layer.requires_grad_(False)
|
||||
#
|
||||
for p in self.get_trainable_params().values():
|
||||
p.requires_grad_(True)
|
||||
|
||||
def state_dict(self):
|
||||
return {n: p.data.cpu() for n, p in self.get_trainable_params().items()}
|
||||
|
||||
def load_state_dict(self, state_dict: Dict[str, torch.Tensor]):
|
||||
params = self.get_trainable_params()
|
||||
for n, p in params.items():
|
||||
p.data[:] = state_dict[n].to(p.data.device)
|
||||
|
||||
for n in state_dict.keys():
|
||||
assert n in params, n
|
||||
|
||||
|
||||
class FineTuneBiases(FineTuner):
|
||||
def get_layer_trainable_params(self, layer: NeoXModule, prefix: str) -> Dict[str, nn.Parameter]:
|
||||
params = {}
|
||||
|
||||
if isinstance(layer, TransformerLayer):
|
||||
# No need to train the mlp bias because we are adding it with attention output
|
||||
params[f'{prefix}.attention.output.bias'] = layer.attention.output.bias
|
||||
params[f'{prefix}.attention.qkv_lin.bias'] = layer.attention.qkv_lin.bias
|
||||
params[f'{prefix}.ffn.dense_h_h4.bias'] = layer.ffn.dense_h_h4.bias
|
||||
else:
|
||||
pass
|
||||
|
||||
return params
|
||||
@@ -0,0 +1,75 @@
|
||||
"""
|
||||
---
|
||||
title: LLM.int8() on GPT-NeoX
|
||||
summary: >
|
||||
Transform nn.Linear layers to 8-bit integer layers.
|
||||
---
|
||||
|
||||
# LLM.int() on GPT-NeoX
|
||||
|
||||
This implements a utility function to transform a `nn.Linear` layer to LLM.int8() linear layer.
|
||||
|
||||
[LLM.int8() paper](https://arxiv.org/abs/eb2bcaee1d0011edaa66a71c10a887e7)
|
||||
shows you can use int8 quantization while handling outliers to
|
||||
reduce memory footprint without performance degradation in large language models.
|
||||
They convert weights and inputs to scaled 8-bit integers and does matrix multiplication
|
||||
producing int32 results which is then converted back to float16 and rescaled.
|
||||
They show that in large langauge models, some features can give extreme values (outliers)
|
||||
that dominate the model's output.
|
||||
These features get clamped in 8-bit integer space which causes the model performance to degrade.
|
||||
As a solution they pick these outliers (greater than a specified threshold)
|
||||
and compute their multiplications separately in float16 space.
|
||||
Since the percentage of outliers is around 0.01% this doesn't increase memory usage,
|
||||
and prevents the model from degrading performance.
|
||||
|
||||
The code to transform GPT-NoeX layers is defined in [model.py](../model.html#post_load_prepare).
|
||||
|
||||
Here are example uses of GPT-NeoX with int8 quantization.
|
||||
|
||||
* [Generate Text](../samples/llm_int8.html)
|
||||
* [Run Evaluation Tests](../evaluation/llm_int8.html)
|
||||
"""
|
||||
|
||||
# Import [`bitsandbytes`](https://github.com/timdettmers/bitsandbytes) package
|
||||
try:
|
||||
from bitsandbytes.nn import Linear8bitLt, Int8Params
|
||||
except ImportError:
|
||||
raise ImportError('''Please install `bitsandbytes` with `pip install bitsandbytes -U`''')
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
|
||||
def make_llm_int8_linear(linear_module: nn.Linear, device: torch.device, threshold: float = 6.0):
|
||||
"""
|
||||
## Transform a `nn.Linear` layer to LLM.int8() linear layer
|
||||
|
||||
:param linear_module: is the `nn.Linear` layer to transform
|
||||
:param device: is the device of the model
|
||||
:param threshold: is the threshold $\alpha$ to use for outlier detection
|
||||
"""
|
||||
|
||||
#
|
||||
assert isinstance(linear_module, nn.Linear)
|
||||
|
||||
# Create an empty Linear8bitLt module
|
||||
int8_lin = Linear8bitLt(
|
||||
linear_module.in_features,
|
||||
linear_module.out_features,
|
||||
linear_module.bias is not None,
|
||||
has_fp16_weights=False,
|
||||
threshold=threshold,
|
||||
)
|
||||
|
||||
# Quantize the weights
|
||||
int8_lin._parameters['weight'] = Int8Params(linear_module.weight.data.cpu(),
|
||||
requires_grad=False,
|
||||
has_fp16_weights=False).to(device)
|
||||
|
||||
# Set the bias in float16 space
|
||||
if linear_module.bias is not None:
|
||||
int8_lin._parameters['bias'] = nn.Parameter(linear_module.bias.data,
|
||||
requires_grad=False)
|
||||
|
||||
#
|
||||
return int8_lin
|
||||
@@ -0,0 +1,132 @@
|
||||
"""
|
||||
---
|
||||
title: Text Dataset for GPT-NeoX
|
||||
summary: >
|
||||
Loads text datasets to fine-tune GPT-NeoX
|
||||
---
|
||||
|
||||
# Text Dataset for GPT-NeoX
|
||||
"""
|
||||
from pathlib import PurePath, Path
|
||||
from typing import Optional, List
|
||||
|
||||
import torch
|
||||
import torch.utils.data
|
||||
from labml import lab
|
||||
from labml import monit
|
||||
from labml.logger import inspect
|
||||
from labml.utils.download import download_file
|
||||
|
||||
from labml_nn.neox.tokenizer import get_tokenizer
|
||||
|
||||
|
||||
def load_text(path: PurePath, url: Optional[str] = None, *, filter_subset: Optional[int] = None):
|
||||
"""
|
||||
### Load text file
|
||||
|
||||
:param path: is the location of the text file
|
||||
:param url: is the URL to download the file from
|
||||
:param filter_subset: is the number of characters to filter.
|
||||
Use this during testing when trying large datasets
|
||||
:return: the text content
|
||||
"""
|
||||
|
||||
path = Path(path)
|
||||
|
||||
# Download if it doesn't exist
|
||||
if not path.exists():
|
||||
if not url:
|
||||
raise FileNotFoundError(str(path))
|
||||
else:
|
||||
download_file(url, path)
|
||||
|
||||
with monit.section("Load data"):
|
||||
# Load data
|
||||
with open(str(path), 'r') as f:
|
||||
text = f.read()
|
||||
# Filter
|
||||
if filter_subset:
|
||||
text = text[:filter_subset]
|
||||
|
||||
#
|
||||
return text
|
||||
|
||||
|
||||
class NeoXDataset(torch.utils.data.Dataset):
|
||||
"""
|
||||
## Dataset for fine-tuning GPT-NeoX
|
||||
|
||||
This is not optimized to very large datasets.
|
||||
"""
|
||||
|
||||
def __init__(self, tokens: List[int], seq_len: int):
|
||||
"""
|
||||
:param tokens: is the list of token ids
|
||||
:param seq_len: is the sequence length of a single training sample
|
||||
"""
|
||||
|
||||
self.seq_len = seq_len
|
||||
# Number of samples
|
||||
n_samples = len(tokens) // seq_len
|
||||
self.n_samples = n_samples
|
||||
# Truncate
|
||||
tokens = tokens[:n_samples * seq_len + 1]
|
||||
# Create a PyTorch tensor
|
||||
self.tokens = torch.tensor(tokens)
|
||||
|
||||
def __len__(self):
|
||||
return self.n_samples
|
||||
|
||||
def __getitem__(self, idx: int):
|
||||
"""
|
||||
### Get a sample
|
||||
|
||||
:param idx: is the index of the sample
|
||||
:return: the input and the target
|
||||
"""
|
||||
offset = idx * self.seq_len
|
||||
return self.tokens[offset:offset + self.seq_len], self.tokens[offset + 1:offset + 1 + self.seq_len]
|
||||
|
||||
|
||||
DATASETS = {
|
||||
'tiny_shakespeare': {
|
||||
'file': 'tiny_shakespeare.txt',
|
||||
'url': 'https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt'
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def get_training_data(seq_len: int = 32, dataset_name: str = 'tiny_shakespeare', truncate: int = -1):
|
||||
"""
|
||||
### Load Dataset
|
||||
|
||||
:param seq_len: is the sequence length of a single training sample
|
||||
:param dataset_name: is the name of the dataset
|
||||
:return: the dataset
|
||||
"""
|
||||
|
||||
ds = DATASETS[dataset_name]
|
||||
# Load the content
|
||||
text = load_text(lab.get_data_path() / ds['file'], ds['url'])
|
||||
# Tokenize
|
||||
tokenizer = get_tokenizer()
|
||||
tokens = tokenizer.encode_batch([text])[0]
|
||||
|
||||
if truncate > 0:
|
||||
token_ids = tokens.ids[:truncate * seq_len]
|
||||
else:
|
||||
token_ids = tokens.ids
|
||||
|
||||
#
|
||||
return NeoXDataset(token_ids, seq_len)
|
||||
|
||||
|
||||
def _test():
|
||||
dataset = get_training_data()
|
||||
|
||||
inspect(tokens=len(dataset.tokens))
|
||||
|
||||
|
||||
#
|
||||
if __name__ == '__main__':
|
||||
_test()
|
||||
@@ -0,0 +1,182 @@
|
||||
from typing import Optional, Set, List
|
||||
|
||||
import torch.nn as nn
|
||||
import torch.optim
|
||||
import torch.utils.data
|
||||
from torch.cuda import amp
|
||||
from torch.cuda.amp import GradScaler
|
||||
|
||||
from labml import monit, tracker
|
||||
from labml.configs import BaseConfigs, option
|
||||
from labml_nn.neox.utils.finetune import FineTuner
|
||||
|
||||
|
||||
def get_trainable_params(model: nn.Module):
|
||||
"""
|
||||
### Get trainable parameters
|
||||
|
||||
:param model: is the model to train
|
||||
:return: a list of parameters for training
|
||||
"""
|
||||
|
||||
# Get all parameters
|
||||
params = list(model.parameters())
|
||||
# Filter parameters that require gradients
|
||||
trainable_params = [p for p in params if p.requires_grad]
|
||||
|
||||
#
|
||||
return trainable_params
|
||||
|
||||
|
||||
class TrainerConf(BaseConfigs):
|
||||
model: nn.Module
|
||||
layers: List[nn.Module]
|
||||
optimizer: torch.optim.Optimizer = 'Adam'
|
||||
train_loader: torch.utils.data.DataLoader
|
||||
valid_loader: Optional[torch.utils.data.DataLoader] = None,
|
||||
device: torch.device = torch.device('cuda:0')
|
||||
scaler: Optional[GradScaler] = 'Default'
|
||||
is_amp: bool = True
|
||||
dtype: torch.dtype = torch.float16
|
||||
|
||||
is_clone_layers: bool = True
|
||||
|
||||
loss_func: nn.Module = nn.CrossEntropyLoss()
|
||||
checkpoints_per_epoch: int = 0
|
||||
samples_per_epoch: int = 0
|
||||
|
||||
grad_norm: Optional[float] = 1.0
|
||||
learning_rate: float = 3e-4
|
||||
max_seq_len: int = 1024
|
||||
batch_size: int = 64
|
||||
epochs: int = 16
|
||||
|
||||
n_gpus: int = torch.cuda.device_count()
|
||||
|
||||
filter_layers: Optional[Set] = None
|
||||
|
||||
def get_loss(self, sample, dataset_split: str):
|
||||
"""
|
||||
:param dataset_split: train/valid
|
||||
:param sample: is the sample
|
||||
:return: the loss, output and the target
|
||||
"""
|
||||
data, target = sample
|
||||
|
||||
# Forward pass
|
||||
with monit.section('Forward pass'):
|
||||
output = self.model(data.to(self.device))
|
||||
# Move targets to the same device as output
|
||||
target = target.to(output.device)
|
||||
# Calculate loss
|
||||
loss = self.loss_func(output.view(target.numel(), -1), target.view(-1))
|
||||
|
||||
return loss, output, target
|
||||
|
||||
def train(self):
|
||||
for epoch in monit.loop(self.epochs):
|
||||
self.train_epoch()
|
||||
tracker.new_line()
|
||||
|
||||
def sample(self, idx):
|
||||
pass
|
||||
|
||||
def save_checkpoint(self, idx):
|
||||
pass
|
||||
|
||||
def get_iterators(self):
|
||||
# Iterate through the batches
|
||||
iterators = [('train', self.train_loader)]
|
||||
if self.valid_loader is not None:
|
||||
iterators.append(('valid', self.valid_loader))
|
||||
|
||||
if self.samples_per_epoch > 0:
|
||||
iterators.append((self.sample, [i for i in range(self.samples_per_epoch)]))
|
||||
|
||||
if self.checkpoints_per_epoch > 0:
|
||||
iterators.append((self.save_checkpoint, [i for i in range(self.checkpoints_per_epoch)]))
|
||||
|
||||
return iterators
|
||||
|
||||
def train_epoch(self):
|
||||
# Set model for train
|
||||
self.model.train()
|
||||
|
||||
iterators = self.get_iterators()
|
||||
for split_name, sample in monit.mix(1024, *iterators):
|
||||
if split_name == 'train':
|
||||
# Set gradients to zero
|
||||
self.optimizer.zero_grad()
|
||||
tracker.add_global_step()
|
||||
|
||||
with torch.set_grad_enabled(split_name == 'train'):
|
||||
if self.is_amp:
|
||||
# Forward pass
|
||||
with amp.autocast():
|
||||
loss, output, target = self.get_loss(sample, split_name)
|
||||
else:
|
||||
loss, output, target = self.get_loss(sample, split_name)
|
||||
|
||||
# Get predictions
|
||||
pred = output.argmax(dim=-1)
|
||||
# Calculate accuracy
|
||||
accuracy = pred.eq(target).sum().item() / (target != -100).sum()
|
||||
|
||||
tracker.add({f'loss.{split_name}': loss, f'acc.{split_name}': accuracy * 100})
|
||||
|
||||
if split_name == 'train':
|
||||
if self.scaler is not None:
|
||||
# Backward pass
|
||||
loss = self.scaler.scale(loss)
|
||||
# tracker.add({'loss.scaled': loss})
|
||||
|
||||
with monit.section('Backward pass'):
|
||||
loss.backward()
|
||||
|
||||
# Optimize
|
||||
with monit.section('Optimize'):
|
||||
if self.scaler is None:
|
||||
self.optimizer.step()
|
||||
else:
|
||||
self.scaler.unscale_(self.optimizer)
|
||||
if self.grad_norm is not None:
|
||||
torch.nn.utils.clip_grad_norm_(get_trainable_params(self.model), self.grad_norm)
|
||||
self.scaler.step(self.optimizer)
|
||||
self.scaler.update()
|
||||
|
||||
tracker.save()
|
||||
|
||||
|
||||
@option(TrainerConf.optimizer, 'Adam')
|
||||
def adam_optimizer(c: TrainerConf):
|
||||
if c.dtype == torch.float32:
|
||||
return torch.optim.Adam(get_trainable_params(c.model), lr=c.learning_rate)
|
||||
elif c.dtype == torch.float16:
|
||||
from labml_nn.optimizers.adam_fp16 import AdamFP16
|
||||
return AdamFP16(get_trainable_params(c.model), lr=c.learning_rate)
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
@option(TrainerConf.optimizer, 'SGD')
|
||||
def sgd_optimizer(c: TrainerConf):
|
||||
return torch.optim.SGD(get_trainable_params(c.model), lr=c.learning_rate)
|
||||
|
||||
|
||||
@option(TrainerConf.scaler, 'Default')
|
||||
def grad_scaler(c: TrainerConf):
|
||||
if not c.is_amp:
|
||||
return None
|
||||
|
||||
if c.dtype == torch.float16:
|
||||
from labml_nn.optimizers.adam_fp16 import GradScalerFP16
|
||||
return GradScalerFP16()
|
||||
else:
|
||||
return GradScaler()
|
||||
|
||||
|
||||
class PipelineParallelTrainerConf(TrainerConf):
|
||||
is_checkpointing: bool = False
|
||||
chunks: int
|
||||
|
||||
fine_tuner: FineTuner
|
||||
Reference in New Issue
Block a user