Spaces:
Sleeping
Sleeping
| import math | |
| import struct | |
| import inspect | |
| from dataclasses import dataclass | |
| from typing import Any, Optional, Tuple, List | |
| import numpy as np | |
| import torch | |
| import torch.nn.functional as F | |
| from torch import nn | |
| class ModelArgs: | |
| dim: int = 1024 | |
| n_layers: int = 12 | |
| n_heads: int = 12 | |
| n_kv_heads: Optional[int] = None | |
| vocab_size: int = 8 * 1024 | |
| hidden_dim: Optional[int] = None | |
| multiple_of: int = 256 # MLP hidden layer size will be multiple of | |
| norm_eps: float = 1e-5 | |
| max_seq_len: int = 1024 | |
| dropout: float = 0.0 | |
| class ConditionalInput: | |
| def __init__(self, name, dtype, min_value, max_value): | |
| super().__init__() | |
| self.name = name | |
| self.dtype = dtype | |
| self.min_value = min_value | |
| self.max_value = max_value | |
| class ConditionalModelArgs(ModelArgs): | |
| dim: int = 1024 | |
| n_layers: int = 12 | |
| n_heads: int = 12 | |
| n_kv_heads: Optional[int] = None | |
| vocab_size: int = 8 * 1024 | |
| hidden_dim: Optional[int] = None | |
| multiple_of: int = 256 # MLP hidden layer size will be multiple of | |
| norm_eps: float = 1e-5 | |
| max_seq_len: int = 1024 | |
| dropout: float = 0.0 | |
| conditions: Optional[List[ConditionalInput]] = None | |
| class KVCache(nn.Module): | |
| def __init__(self, max_batch_size, max_seq_length, n_head, head_dim, dtype): | |
| super().__init__() | |
| cache_shape = (max_batch_size, n_head, max_seq_length, head_dim) | |
| self.register_buffer('k_cache', torch.zeros(cache_shape, dtype=dtype)) | |
| self.register_buffer('v_cache', torch.zeros(cache_shape, dtype=dtype)) | |
| def update(self, input_pos, k_val, v_val): | |
| # input_pos: [S], k_val: [B, H, S, D] | |
| assert input_pos.shape[0] == k_val.shape[2] | |
| k_out = self.k_cache | |
| v_out = self.v_cache | |
| k_out[:, :, input_pos] = k_val | |
| v_out[:, :, input_pos] = v_val | |
| return k_out, v_out | |
| class RMSNorm(nn.Module): | |
| def __init__(self, dim: int, eps: float): | |
| super().__init__() | |
| self.eps = eps | |
| self.weight = nn.Parameter(torch.ones(dim)) | |
| def _norm(self, x): | |
| return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) | |
| def forward(self, x): | |
| output = self._norm(x.float()).type_as(x) | |
| return output * self.weight | |
| def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0): | |
| freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)) | |
| t = torch.arange(end, device=freqs.device) # type: ignore | |
| freqs = torch.outer(t, freqs).float() # type: ignore | |
| freqs_cos = torch.cos(freqs) # real part | |
| freqs_sin = torch.sin(freqs) # imaginary part | |
| return freqs_cos, freqs_sin | |
| def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor): | |
| ndim = x.ndim | |
| assert 0 <= 1 < ndim | |
| assert freqs_cis.shape == (x.shape[1], x.shape[-1]) | |
| shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)] | |
| return freqs_cis.view(shape) | |
| def apply_rotary_emb( | |
| xq: torch.Tensor, | |
| xk: torch.Tensor, | |
| freqs_cos: torch.Tensor, | |
| freqs_sin: torch.Tensor | |
| ) -> Tuple[torch.Tensor, torch.Tensor]: | |
| # reshape xq and xk to match the complex representation | |
| xq_r, xq_i = xq.float().reshape(xq.shape[:-1] + (-1, 2)).unbind(-1) | |
| xk_r, xk_i = xk.float().reshape(xk.shape[:-1] + (-1, 2)).unbind(-1) | |
| # reshape freqs_cos and freqs_sin for broadcasting | |
| freqs_cos = reshape_for_broadcast(freqs_cos, xq_r) | |
| freqs_sin = reshape_for_broadcast(freqs_sin, xq_r) | |
| # apply rotation using real numbers | |
| xq_out_r = xq_r * freqs_cos - xq_i * freqs_sin | |
| xq_out_i = xq_r * freqs_sin + xq_i * freqs_cos | |
| xk_out_r = xk_r * freqs_cos - xk_i * freqs_sin | |
| xk_out_i = xk_r * freqs_sin + xk_i * freqs_cos | |
| # flatten last two dimensions | |
| xq_out = torch.stack([xq_out_r, xq_out_i], dim=-1).flatten(3) | |
| xk_out = torch.stack([xk_out_r, xk_out_i], dim=-1).flatten(3) | |
| return xq_out.type_as(xq), xk_out.type_as(xk) | |
| def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor: | |
| """torch.repeat_interleave(x, dim=2, repeats=n_rep)""" | |
| bs, slen, n_kv_heads, head_dim = x.shape | |
| if n_rep == 1: | |
| return x | |
| return ( | |
| x[:, :, :, None, :] | |
| .expand(bs, slen, n_kv_heads, n_rep, head_dim) | |
| .reshape(bs, slen, n_kv_heads * n_rep, head_dim) | |
| ) | |
| class Attention(nn.Module): | |
| def __init__(self, args: ModelArgs): | |
| super().__init__() | |
| self.n_kv_heads = args.n_heads if args.n_kv_heads is None else args.n_kv_heads | |
| assert args.n_heads % self.n_kv_heads == 0 | |
| model_parallel_size = 1 | |
| self.n_local_heads = args.n_heads // model_parallel_size | |
| self.n_local_kv_heads = self.n_kv_heads // model_parallel_size | |
| self.n_rep = self.n_local_heads // self.n_local_kv_heads | |
| self.head_dim = args.dim // args.n_heads | |
| self.wq = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False) | |
| self.wk = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False) | |
| self.wv = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False) | |
| self.wo = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False) | |
| self.attn_dropout = nn.Dropout(args.dropout) | |
| self.resid_dropout = nn.Dropout(args.dropout) | |
| self.dropout = args.dropout | |
| # use flash attention or a manual implementation? | |
| self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') | |
| if not self.flash: | |
| print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0") | |
| mask = torch.full((1, 1, args.max_seq_len, args.max_seq_len), float("-inf")) | |
| mask = torch.triu(mask, diagonal=1) | |
| self.register_buffer("mask", mask) | |
| def forward( | |
| self, | |
| x: torch.Tensor, | |
| freqs_cos: torch.Tensor, | |
| freqs_sin: torch.Tensor, | |
| ): | |
| bsz, seqlen, _ = x.shape | |
| # QKV | |
| xq, xk, xv = self.wq(x), self.wk(x), self.wv(x) | |
| xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim) | |
| xk = xk.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim) | |
| xv = xv.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim) | |
| # RoPE relative positional embeddings | |
| xq, xk = apply_rotary_emb(xq, xk, freqs_cos, freqs_sin) | |
| # grouped multiquery attention: expand out keys and values | |
| xk = repeat_kv(xk, self.n_rep) # (bs, seqlen, n_local_heads, head_dim) | |
| xv = repeat_kv(xv, self.n_rep) # (bs, seqlen, n_local_heads, head_dim) | |
| # make heads into a batch dimension | |
| xq = xq.transpose(1, 2) # (bs, n_local_heads, seqlen, head_dim) | |
| xk = xk.transpose(1, 2) | |
| xv = xv.transpose(1, 2) | |
| # flash implementation | |
| if self.flash: | |
| output = torch.nn.functional.scaled_dot_product_attention(xq, xk, xv, attn_mask=None, dropout_p=self.dropout if self.training else 0.0, is_causal=True) | |
| else: | |
| # manual implementation | |
| scores = torch.matmul(xq, xk.transpose(2, 3)) / math.sqrt(self.head_dim) | |
| assert hasattr(self, 'mask') | |
| scores = scores + self.mask[:, :, :seqlen, :seqlen] # (bs, n_local_heads, seqlen, cache_len + seqlen) | |
| scores = F.softmax(scores.float(), dim=-1).type_as(xq) | |
| scores = self.attn_dropout(scores) | |
| output = torch.matmul(scores, xv) # (bs, n_local_heads, seqlen, head_dim) | |
| # restore time as batch dimension and concat heads | |
| output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1) | |
| # final projection into the residual stream | |
| output = self.wo(output) | |
| output = self.resid_dropout(output) | |
| return output | |
| class FeedForward(nn.Module): | |
| def __init__(self, dim: int, hidden_dim: int, multiple_of: int, dropout: float): | |
| super().__init__() | |
| if hidden_dim is None: | |
| hidden_dim = 4 * dim | |
| hidden_dim = int(2 * hidden_dim / 3) | |
| hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of) | |
| self.w1 = nn.Linear(dim, hidden_dim, bias=False) | |
| self.w2 = nn.Linear(hidden_dim, dim, bias=False) | |
| self.w3 = nn.Linear(dim, hidden_dim, bias=False) | |
| self.dropout = nn.Dropout(dropout) | |
| def forward(self, x): | |
| return self.dropout(self.w2(F.silu(self.w1(x)) * self.w3(x))) | |
| class TransformerBlock(nn.Module): | |
| def __init__(self, layer_id: int, args: ModelArgs): | |
| super().__init__() | |
| self.n_heads = args.n_heads | |
| self.dim = args.dim | |
| self.head_dim = args.dim // args.n_heads | |
| self.attention = Attention(args) | |
| self.feed_forward = FeedForward( | |
| dim=args.dim, | |
| hidden_dim=args.hidden_dim, | |
| multiple_of=args.multiple_of, | |
| dropout=args.dropout, | |
| ) | |
| self.layer_id = layer_id | |
| self.attention_norm = RMSNorm(args.dim, eps=args.norm_eps) | |
| self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps) | |
| def forward(self, x, freqs_cos, freqs_sin): | |
| h = x + self.attention.forward(self.attention_norm(x), freqs_cos, freqs_sin) | |
| out = h + self.feed_forward.forward(self.ffn_norm(h)) | |
| return out | |
| class Transformer(nn.Module): | |
| last_loss: Optional[torch.Tensor] | |
| def __init__(self, params: ModelArgs): | |
| super().__init__() | |
| self.params = params | |
| self.vocab_size = params.vocab_size | |
| self.n_layers = params.n_layers | |
| self.tok_embeddings = nn.Embedding(params.vocab_size, params.dim) | |
| self.dropout = nn.Dropout(params.dropout) | |
| self.layers = torch.nn.ModuleList() | |
| for layer_id in range(params.n_layers): | |
| self.layers.append(TransformerBlock(layer_id, params)) | |
| self.norm = RMSNorm(params.dim, eps=params.norm_eps) | |
| self.output = nn.Linear(params.dim, params.vocab_size, bias=False) | |
| # share the unembedding parameters with the embedding parameters | |
| self.tok_embeddings.weight = self.output.weight # https://paperswithcode.com/method/weight-tying | |
| # some useful precompute for the RoPE relative positional embeddings | |
| freqs_cos, freqs_sin = precompute_freqs_cis(self.params.dim // self.params.n_heads, self.params.max_seq_len) | |
| self.register_buffer("freqs_cos", freqs_cos, persistent=False) | |
| self.register_buffer("freqs_sin", freqs_sin, persistent=False) | |
| # init all weights | |
| self.apply(self._init_weights) | |
| # apply special scaled init to the residual projections, per GPT-2 paper | |
| for pn, p in self.named_parameters(): | |
| if pn.endswith('w3.weight') or pn.endswith('wo.weight'): | |
| torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * params.n_layers)) | |
| # Initialize attribute for the loss of the last forward call. This will be set if the forward is called with a targets tensor. | |
| self.last_loss = None | |
| def _init_weights(self, module): | |
| if isinstance(module, nn.Linear): | |
| torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) | |
| if module.bias is not None: | |
| torch.nn.init.zeros_(module.bias) | |
| elif isinstance(module, nn.Embedding): | |
| torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) | |
| def forward(self, tokens: torch.Tensor, targets: Optional[torch.Tensor] = None, full_logits=False) -> torch.Tensor: | |
| _bsz, seqlen = tokens.shape | |
| h = self.tok_embeddings(tokens) | |
| h = self.dropout(h) | |
| freqs_cos = self.freqs_cos[:seqlen] | |
| freqs_sin = self.freqs_sin[:seqlen] | |
| for layer in self.layers: | |
| h = layer(h, freqs_cos, freqs_sin) | |
| h = self.norm(h) | |
| if targets is not None: | |
| # if we are given some desired targets also calculate the loss | |
| logits = self.output(h) | |
| self.last_loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.reshape(-1), ignore_index=-1) | |
| elif full_logits: | |
| logits = self.output(h) | |
| else: | |
| # inference-time mini-optimization: only forward the output on the very last position | |
| logits = self.output(h[:, [-1], :]) # note: using list [-1] to preserve the time dim | |
| self.last_loss = None | |
| return logits | |
| def configure_optimizers(self, weight_decay, learning_rate, betas, device_type): | |
| # start with all of the candidate parameters | |
| param_dict = {pn: p for pn, p in self.named_parameters()} | |
| # filter out those that do not require grad | |
| param_dict = {pn: p for pn, p in param_dict.items() if p.requires_grad} | |
| # create optim groups. Any parameters that is 2D will be weight decayed, otherwise no. | |
| # i.e. all weight tensors in matmuls + embeddings decay, all biases and layernorms don't. | |
| decay_params = [p for n, p in param_dict.items() if p.dim() >= 2] | |
| nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2] | |
| optim_groups = [ | |
| {'params': decay_params, 'weight_decay': weight_decay}, | |
| {'params': nodecay_params, 'weight_decay': 0.0} | |
| ] | |
| num_decay_params = sum(p.numel() for p in decay_params) | |
| num_nodecay_params = sum(p.numel() for p in nodecay_params) | |
| print(f"num decayed parameter tensors: {len(decay_params)}, with {num_decay_params:,} parameters") | |
| print(f"num non-decayed parameter tensors: {len(nodecay_params)}, with {num_nodecay_params:,} parameters") | |
| # Create AdamW optimizer and use the fused version if it is available | |
| fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters | |
| use_fused = fused_available and device_type == 'cuda' | |
| extra_args = dict(fused=True) if use_fused else dict() | |
| optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas, **extra_args) | |
| print(f"using fused AdamW: {use_fused}") | |
| return optimizer | |
| def estimate_mfu(self, fwdbwd_per_iter, dt): | |
| """ estimate model flops utilization (MFU) in units of A100 bfloat16 peak FLOPS """ | |
| # first estimate the number of flops we do per iteration. | |
| # see PaLM paper Appendix B as ref: https://arxiv.org/abs/2204.02311 | |
| N = sum(p.numel() for p in self.parameters()) | |
| cfg = self.params | |
| L, H, Q, T = cfg.n_layers, cfg.n_heads, cfg.dim//cfg.n_heads, cfg.max_seq_len | |
| flops_per_token = 6*N + 12*L*H*Q*T | |
| flops_per_fwdbwd = flops_per_token * T | |
| flops_per_iter = flops_per_fwdbwd * fwdbwd_per_iter | |
| # express our flops throughput as ratio of A100 bfloat16 peak flops | |
| flops_achieved = flops_per_iter * (1.0/dt) # per second | |
| flops_promised = 312e12 # A100 GPU bfloat16 peak flops is 312 TFLOPS | |
| mfu = flops_achieved / flops_promised | |
| return mfu | |
| def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None): | |
| """ | |
| Take a conditioning sequence of indices idx (LongTensor of shape (b,t)) and complete | |
| the sequence max_new_tokens times, feeding the predictions back into the model each time. | |
| Most likely you'll want to make sure to be in model.eval() mode of operation for this. | |
| Also note this is a super inefficient version of sampling with no key/value cache. | |
| """ | |
| for _ in range(max_new_tokens): | |
| # if the sequence context is growing too long we must crop it at block_size | |
| idx_cond = idx if idx.size(1) <= self.params.max_seq_len else idx[:, -self.params.max_seq_len:] | |
| # forward the model to get the logits for the index in the sequence | |
| logits = self(idx_cond) | |
| logits = logits[:, -1, :] # crop to just the final time step | |
| if temperature == 0.0: | |
| # "sample" the single most likely index | |
| _, idx_next = torch.topk(logits, k=1, dim=-1) | |
| else: | |
| # pluck the logits at the final step and scale by desired temperature | |
| logits = logits / temperature | |
| # optionally crop the logits to only the top k options | |
| if top_k is not None: | |
| v, _ = torch.topk(logits, min(top_k, logits.size(-1))) | |
| logits[logits < v[:, [-1]]] = -float('Inf') | |
| # apply softmax to convert logits to (normalized) probabilities | |
| probs = F.softmax(logits, dim=-1) | |
| idx_next = torch.multinomial(probs, num_samples=1) | |
| # append sampled index to the running sequence and continue | |
| idx = torch.cat((idx, idx_next), dim=1) | |
| return idx | |
| class RegressTransformer(nn.Module): | |
| last_loss: Optional[torch.Tensor] | |
| def __init__(self, params: ModelArgs): | |
| super().__init__() | |
| self.params = params | |
| self.vocab_size = params.vocab_size | |
| self.n_layers = params.n_layers | |
| self.tok_embeddings = nn.Embedding(params.vocab_size, params.dim) | |
| self.dropout = nn.Dropout(params.dropout) | |
| self.layers = torch.nn.ModuleList() | |
| for layer_id in range(params.n_layers): | |
| self.layers.append(TransformerBlock(layer_id, params)) | |
| self.norm = RMSNorm(params.dim, eps=params.norm_eps) | |
| self.output = nn.Linear(params.dim, 1)#, bias=False) | |
| # share the unembedding parameters with the embedding parameters | |
| # some useful precompute for the RoPE relative positional embeddings | |
| freqs_cos, freqs_sin = precompute_freqs_cis(self.params.dim // self.params.n_heads, self.params.max_seq_len) | |
| self.register_buffer("freqs_cos", freqs_cos, persistent=False) | |
| self.register_buffer("freqs_sin", freqs_sin, persistent=False) | |
| # init all weights | |
| self.apply(self._init_weights) | |
| # apply special scaled init to the residual projections, per GPT-2 paper | |
| for pn, p in self.named_parameters(): | |
| if pn.endswith('w3.weight') or pn.endswith('wo.weight'): | |
| torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * params.n_layers)) | |
| # Initialize attribute for the loss of the last forward call. This will be set if the forward is called with a targets tensor. | |
| self.last_loss = None | |
| def _init_weights(self, module): | |
| if isinstance(module, nn.Linear): | |
| torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) | |
| if module.bias is not None: | |
| torch.nn.init.zeros_(module.bias) | |
| elif isinstance(module, nn.Embedding): | |
| torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) | |
| def forward(self, tokens: torch.Tensor, targets: Optional[dict] = None, full_logits=False) -> torch.Tensor: | |
| _bsz, seqlen = tokens.shape | |
| h = self.tok_embeddings(tokens) | |
| h = self.dropout(h) | |
| freqs_cos = self.freqs_cos[:seqlen] | |
| freqs_sin = self.freqs_sin[:seqlen] | |
| for layer in self.layers: | |
| h = layer(h, freqs_cos, freqs_sin) | |
| h = self.norm(h) | |
| h = h.mean(dim=1) | |
| if targets is not None: | |
| # if we are given some desired targets also calculate the loss | |
| logits = self.output(h).squeeze(-1) | |
| self.last_loss = F.mse_loss(logits, targets['QED']) | |
| else: | |
| # inference-time mini-optimization: only forward the output on the very last position | |
| # logits = self.output(h[:, [-1], :]) # note: using list [-1] to preserve the time dim | |
| logits = h | |
| self.last_loss = None | |
| return logits | |
| def configure_optimizers(self, weight_decay, learning_rate, betas, device_type): | |
| # start with all of the candidate parameters | |
| param_dict = {pn: p for pn, p in self.named_parameters()} | |
| # filter out those that do not require grad | |
| param_dict = {pn: p for pn, p in param_dict.items() if p.requires_grad} | |
| # create optim groups. Any parameters that is 2D will be weight decayed, otherwise no. | |
| # i.e. all weight tensors in matmuls + embeddings decay, all biases and layernorms don't. | |
| decay_params = [p for n, p in param_dict.items() if p.dim() >= 2] | |
| nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2] | |
| optim_groups = [ | |
| {'params': decay_params, 'weight_decay': weight_decay}, | |
| {'params': nodecay_params, 'weight_decay': 0.0} | |
| ] | |
| num_decay_params = sum(p.numel() for p in decay_params) | |
| num_nodecay_params = sum(p.numel() for p in nodecay_params) | |
| print(f"num decayed parameter tensors: {len(decay_params)}, with {num_decay_params:,} parameters") | |
| print(f"num non-decayed parameter tensors: {len(nodecay_params)}, with {num_nodecay_params:,} parameters") | |
| # Create AdamW optimizer and use the fused version if it is available | |
| fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters | |
| use_fused = fused_available and device_type == 'cuda' | |
| extra_args = dict(fused=True) if use_fused else dict() | |
| optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas, **extra_args) | |
| print(f"using fused AdamW: {use_fused}") | |
| return optimizer | |
| def configure_optimizers_1st_stage(self, weight_decay, learning_rate, betas, device_type): | |
| high_param_dict = {pn: p for pn, p in self.named_parameters() if 'output' in pn} | |
| high_param_dict = {pn: p for pn, p in high_param_dict.items() if p.requires_grad} | |
| low_param_dict = {pn: p for pn, p in self.named_parameters() if pn not in high_param_dict} | |
| low_param_dict = {pn: p for pn, p in low_param_dict.items() if p.requires_grad} | |
| high_decay_params = [p for n, p in high_param_dict.items() if p.dim() >= 2] | |
| high_nodecay_params = [p for n, p in high_param_dict.items() if p.dim() < 2] | |
| low_decay_params = [p for n, p in low_param_dict.items() if p.dim() >= 2] | |
| low_nodecay_params = [p for n, p in low_param_dict.items() if p.dim() < 2] | |
| optim_groups = [ | |
| {'params': high_decay_params, 'weight_decay': weight_decay, 'lr': learning_rate}, | |
| {'params': high_nodecay_params, 'weight_decay': 0.0, 'lr': learning_rate}, | |
| {'params': low_decay_params, 'weight_decay': weight_decay, 'lr': learning_rate / 100}, | |
| {'params': low_nodecay_params, 'weight_decay': 0.0, 'lr': learning_rate / 100}, | |
| ] | |
| # Create AdamW optimizer and use the fused version if it is available | |
| fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters | |
| use_fused = fused_available and device_type == 'cuda' | |
| extra_args = dict(fused=True) if use_fused else dict() | |
| optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas, **extra_args) | |
| print(f"using fused AdamW: {use_fused}") | |
| return optimizer | |
| def configure_optimizers_2nd_stage(self, weight_decay, learning_rate, betas, device_type): | |
| high_param_dict = {pn: p for pn, p in self.named_parameters() if 'output' in pn} | |
| high_param_dict = {pn: p for pn, p in high_param_dict.items() if p.requires_grad} | |
| mid_param_dict = {pn: p for pn, p in self.named_parameters() if 'embedding' in pn} | |
| mid_param_dict = {pn: p for pn, p in mid_param_dict.items() if p.requires_grad} | |
| low_param_dict = {pn: p for pn, p in self.named_parameters() if pn not in high_param_dict and pn not in mid_param_dict} | |
| low_param_dict = {pn: p for pn, p in low_param_dict.items() if p.requires_grad} | |
| high_decay_params = [p for n, p in high_param_dict.items() if p.dim() >= 2] | |
| high_nodecay_params = [p for n, p in high_param_dict.items() if p.dim() < 2] | |
| mid_decay_params = [p for n, p in mid_param_dict.items() if p.dim() >= 2] | |
| mid_nodecay_params = [p for n, p in mid_param_dict.items() if p.dim() < 2] | |
| low_decay_params = [p for n, p in low_param_dict.items() if p.dim() >= 2] | |
| low_nodecay_params = [p for n, p in low_param_dict.items() if p.dim() < 2] | |
| optim_groups = [ | |
| {'params': high_decay_params, 'weight_decay': weight_decay, 'lr': learning_rate}, | |
| {'params': high_nodecay_params, 'weight_decay': 0.0, 'lr': learning_rate}, | |
| {'params': mid_decay_params, 'weight_decay': weight_decay, 'lr': learning_rate / 10}, | |
| {'params': mid_nodecay_params, 'weight_decay': 0.0, 'lr': learning_rate / 10}, | |
| {'params': low_decay_params, 'weight_decay': weight_decay, 'lr': learning_rate / 100}, | |
| {'params': low_nodecay_params, 'weight_decay': 0.0, 'lr': learning_rate / 100}, | |
| ] | |
| # Create AdamW optimizer and use the fused version if it is available | |
| fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters | |
| use_fused = fused_available and device_type == 'cuda' | |
| extra_args = dict(fused=True) if use_fused else dict() | |
| optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas, **extra_args) | |
| print(f"using fused AdamW: {use_fused}") | |
| return optimizer | |
| def estimate_mfu(self, fwdbwd_per_iter, dt): | |
| """ estimate model flops utilization (MFU) in units of A100 bfloat16 peak FLOPS """ | |
| # first estimate the number of flops we do per iteration. | |
| # see PaLM paper Appendix B as ref: https://arxiv.org/abs/2204.02311 | |
| N = sum(p.numel() for p in self.parameters()) | |
| cfg = self.params | |
| L, H, Q, T = cfg.n_layers, cfg.n_heads, cfg.dim//cfg.n_heads, cfg.max_seq_len | |
| flops_per_token = 6*N + 12*L*H*Q*T | |
| flops_per_fwdbwd = flops_per_token * T | |
| flops_per_iter = flops_per_fwdbwd * fwdbwd_per_iter | |
| # express our flops throughput as ratio of A100 bfloat16 peak flops | |
| flops_achieved = flops_per_iter * (1.0/dt) # per second | |
| flops_promised = 312e12 # A100 GPU bfloat16 peak flops is 312 TFLOPS | |
| mfu = flops_achieved / flops_promised | |
| return mfu | |
| class ConditionalEmbedding(nn.Module): | |
| def __init__(self, params: ConditionalModelArgs): | |
| super().__init__() | |
| self.embeddings = nn.ModuleDict() | |
| self.conditions = {} | |
| for condition in params.conditions: | |
| self.conditions[condition.name] = condition | |
| if condition.dtype == 'float': | |
| embedding = nn.Linear(1, params.dim, bias=True) | |
| elif condition.dtype == 'int': | |
| embedding = nn.Embedding(condition.max_value - condition.min_value, params.dim) | |
| else: | |
| raise NotImplementedError(f'dtype {condition.dtype} not implemented in ConditionalEmbedding') | |
| self.embeddings[condition.name] = embedding | |
| def process(self, name, value): | |
| dtype = self.conditions[name].dtype | |
| min_value = self.conditions[name].min_value | |
| max_value = self.conditions[name].max_value | |
| value = torch.clamp(value, min_value, max_value) | |
| if dtype == 'float': | |
| value = (value - min_value) / (max_value - min_value) | |
| elif dtype == 'int': | |
| value = value - min_value | |
| else: | |
| raise NotImplementedError(f'dtype {dtype} not implemented in ConditionalInput') | |
| return value | |
| def forward(self, conditions): | |
| conds = [] | |
| for name, x in conditions.items(): | |
| x = self.process(name, x) | |
| if isinstance(self.embeddings[name], nn.Embedding): | |
| conds.append(self.embeddings[name](x.long())) | |
| elif isinstance(self.embeddings[name], nn.Linear): | |
| conds.append(self.embeddings[name](x.unsqueeze(1).float())) | |
| else: | |
| raise ValueError(f'{type(self.embeddings[name])}, {name}') | |
| if conds: | |
| conds = torch.stack(conds, dim=1) | |
| return conds | |
| return None | |
| use_cfg_embedding = dropout_prob > 0 | |
| self.embedding_table = nn.Embedding(num_classes + use_cfg_embedding, hidden_size) | |
| self.num_classes = num_classes | |
| self.dropout_prob = dropout_prob | |
| def forward(self, labels, train, force_drop_ids=None): | |
| use_dropout = self.dropout_prob > 0 | |
| if (train and use_dropout) or (force_drop_ids is not None): | |
| labels = self.token_drop(labels, force_drop_ids) | |
| embeddings = self.embedding_table(labels).unsqueeze(1) | |
| return embeddings | |
| class SimpleConditionalEmbedding(nn.Module): | |
| def __init__(self, params: ConditionalModelArgs): | |
| super().__init__() | |
| embeddings = {} | |
| null_idxs = {} | |
| for condition in params.conditions: | |
| num_embeddings = condition.max_value - condition.min_value + 1 | |
| null_idxs[condition.name] = num_embeddings | |
| embeddings[condition.name] = nn.Embedding(num_embeddings + 1, params.dim) | |
| self.embeddings = nn.ModuleDict(embeddings) | |
| self.null_idxs = null_idxs | |
| self.individual_dropout_prob = 0.1 | |
| self.group_dropout_prob = 0.2 | |
| def token_drop(self, labels, null_idx, force_drop_ids=None): | |
| """ | |
| Drops labels to enable classifier-free guidance. | |
| """ | |
| if force_drop_ids is None: | |
| if torch.rand(1) < self.group_dropout_prob: | |
| drop_ids = torch.ones(labels.shape[0], device=labels.device).bool() | |
| else: | |
| drop_ids = torch.rand(labels.shape[0], device=labels.device) < self.individual_dropout_prob | |
| else: | |
| drop_ids = force_drop_ids == 1 | |
| labels = torch.where(drop_ids, null_idx, labels) | |
| return labels | |
| def forward(self, x): | |
| conds = [] | |
| for k, v in x.items(): | |
| if self.training: | |
| v = self.token_drop(v, self.null_idxs[k]) | |
| cond = self.embeddings[k](v) | |
| conds.append(cond) | |
| conds = torch.stack(conds, dim=1) | |
| return conds | |
| class ContinuousConditionalEmbedding(nn.Module): | |
| def __init__(self, params: ConditionalModelArgs): | |
| super().__init__() | |
| embeddings = {} | |
| for condition in params.conditions: | |
| embeddings[condition.name] = nn.Sequential(nn.Linear(1, params.dim, bias=False), nn.Dropout(0)) | |
| # embeddings[condition.name] = nn.Sequential( | |
| # nn.Linear(1, params.dim), nn.SiLU(), nn.Linear(params.dim, params.dim), nn.Dropout(0) | |
| # ) | |
| self.embeddings = nn.ModuleDict(embeddings) | |
| self.null_embs = nn.Embedding(1, params.dim) | |
| self.individual_dropout_prob = 0#0.1 | |
| self.group_dropout_prob = 0#.2 | |
| def token_drop(self, labels, null_idx, force_drop_ids=None): | |
| """ | |
| Drops labels to enable classifier-free guidance. | |
| """ | |
| if force_drop_ids is None: | |
| if torch.rand(1) < self.group_dropout_prob: | |
| drop_ids = torch.ones(labels.shape[0], device=labels.device).bool() | |
| else: | |
| # will drop all instances of that condiitoning from batch, not as diverse as dropping over (B, C) | |
| drop_ids = torch.rand(labels.shape[0], device=labels.device) < self.individual_dropout_prob | |
| else: | |
| drop_ids = force_drop_ids == 1 | |
| labels = torch.where(drop_ids, null_idx, labels) | |
| return labels | |
| def forward(self, x): | |
| conds = [] | |
| for k, v in x.items(): | |
| v = self.embeddings[k](v.unsqueeze(-1)) | |
| # if self.training: | |
| # v = self.token_drop(v, 0) | |
| # v = self.null_embs(v.long()) | |
| # else: | |
| # v = self.embeddings[k](v.unsqueeze(-1))# + torch.clamp(0.02 * torch.randn((v.shape[0], 1), device=v.device), 0, 1)) | |
| conds.append(v) | |
| conds = torch.stack(conds, dim=1) | |
| return conds | |
| class ConditionalTransformer(nn.Module): | |
| last_loss: Optional[torch.Tensor] | |
| def __init__(self, params: ModelArgs): | |
| super().__init__() | |
| self.params = params | |
| self.vocab_size = params.vocab_size | |
| self.n_layers = params.n_layers | |
| self.tok_embeddings = nn.Embedding(params.vocab_size, params.dim) | |
| self.dropout = nn.Dropout(params.dropout) | |
| self.layers = torch.nn.ModuleList() | |
| for layer_id in range(params.n_layers): | |
| self.layers.append(TransformerBlock(layer_id, params)) | |
| self.norm = RMSNorm(params.dim, eps=params.norm_eps) | |
| self.output = nn.Linear(params.dim, params.vocab_size, bias=False) | |
| # share the unembedding parameters with the embedding parameters | |
| # self.tok_embeddings.weight = self.output.weight # https://paperswithcode.com/method/weight-tying | |
| # some useful precompute for the RoPE relative positional embeddings | |
| freqs_cos, freqs_sin = precompute_freqs_cis(self.params.dim // self.params.n_heads, self.params.max_seq_len) | |
| self.register_buffer("freqs_cos", freqs_cos, persistent=False) | |
| self.register_buffer("freqs_sin", freqs_sin, persistent=False) | |
| # init all weights | |
| self.apply(self._init_weights) | |
| # apply special scaled init to the residual projections, per GPT-2 paper | |
| for pn, p in self.named_parameters(): | |
| if pn.endswith('w3.weight') or pn.endswith('wo.weight'): | |
| torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * params.n_layers)) | |
| # Initialize attribute for the loss of the last forward call. This will be set if the forward is called with a targets tensor. | |
| self.last_loss = None | |
| def _init_weights(self, module): | |
| if isinstance(module, nn.Linear): | |
| torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) | |
| if module.bias is not None: | |
| torch.nn.init.zeros_(module.bias) | |
| elif isinstance(module, nn.Embedding): | |
| torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) | |
| def forward(self, tokens: torch.Tensor, targets: Optional[torch.Tensor] = None, full_logits=False) -> torch.Tensor: | |
| _bsz, seqlen = tokens.shape | |
| h = self.tok_embeddings(tokens) | |
| h = self.dropout(h) | |
| freqs_cos = self.freqs_cos[:seqlen] | |
| freqs_sin = self.freqs_sin[:seqlen] | |
| for layer in self.layers: | |
| h = layer(h, freqs_cos, freqs_sin) | |
| h = self.norm(h) | |
| if targets is not None: | |
| # if we are given some desired targets also calculate the loss | |
| logits = self.output(h) | |
| loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.reshape(-1), ignore_index=-1, reduction='none') | |
| loss = loss[targets.reshape(-1) < 4096] # dont compute loss over conditioning tokens | |
| self.last_loss = loss.mean() | |
| elif full_logits: | |
| logits = self.output(h) | |
| else: | |
| # inference-time mini-optimization: only forward the output on the very last position | |
| logits = self.output(h[:, [-1], :]) # note: using list [-1] to preserve the time dim | |
| self.last_loss = None | |
| return logits | |
| def configure_optimizers(self, weight_decay, learning_rate, betas, device_type): | |
| # start with all of the candidate parameters | |
| param_dict = {pn: p for pn, p in self.named_parameters()} | |
| # filter out those that do not require grad | |
| param_dict = {pn: p for pn, p in param_dict.items() if p.requires_grad} | |
| # create optim groups. Any parameters that is 2D will be weight decayed, otherwise no. | |
| # i.e. all weight tensors in matmuls + embeddings decay, all biases and layernorms don't. | |
| decay_params = [p for n, p in param_dict.items() if p.dim() >= 2] | |
| nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2] | |
| optim_groups = [ | |
| {'params': decay_params, 'weight_decay': weight_decay}, | |
| {'params': nodecay_params, 'weight_decay': 0.0} | |
| ] | |
| num_decay_params = sum(p.numel() for p in decay_params) | |
| num_nodecay_params = sum(p.numel() for p in nodecay_params) | |
| print(f"num decayed parameter tensors: {len(decay_params)}, with {num_decay_params:,} parameters") | |
| print(f"num non-decayed parameter tensors: {len(nodecay_params)}, with {num_nodecay_params:,} parameters") | |
| # Create AdamW optimizer and use the fused version if it is available | |
| fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters | |
| use_fused = fused_available and device_type == 'cuda' | |
| extra_args = dict(fused=True) if use_fused else dict() | |
| optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas, **extra_args) | |
| print(f"using fused AdamW: {use_fused}") | |
| return optimizer | |
| def configure_optimizers_1st_stage(self, weight_decay, learning_rate, betas, device_type): | |
| high_param_dict = {pn: p for pn, p in self.named_parameters() if 'tok_embeddings' in pn or 'output' in pn} | |
| high_param_dict = {pn: p for pn, p in high_param_dict.items() if p.requires_grad} | |
| low_param_dict = {pn: p for pn, p in self.named_parameters() if pn not in high_param_dict} | |
| low_param_dict = {pn: p for pn, p in low_param_dict.items() if p.requires_grad} | |
| high_decay_params = [p for n, p in high_param_dict.items() if p.dim() >= 2] | |
| high_nodecay_params = [p for n, p in high_param_dict.items() if p.dim() < 2] | |
| low_decay_params = [p for n, p in low_param_dict.items() if p.dim() >= 2] | |
| low_nodecay_params = [p for n, p in low_param_dict.items() if p.dim() < 2] | |
| optim_groups = [ | |
| {'params': high_decay_params, 'weight_decay': weight_decay, 'lr': learning_rate}, | |
| {'params': high_nodecay_params, 'weight_decay': 0.0, 'lr': learning_rate}, | |
| {'params': low_decay_params, 'weight_decay': weight_decay, 'lr': learning_rate / 100}, | |
| {'params': low_nodecay_params, 'weight_decay': 0.0, 'lr': learning_rate / 100}, | |
| ] | |
| # Create AdamW optimizer and use the fused version if it is available | |
| fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters | |
| use_fused = fused_available and device_type == 'cuda' | |
| extra_args = dict(fused=True) if use_fused else dict() | |
| optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas, **extra_args) | |
| print(f"using fused AdamW: {use_fused}") | |
| return optimizer | |
| def estimate_mfu(self, fwdbwd_per_iter, dt): | |
| """ estimate model flops utilization (MFU) in units of A100 bfloat16 peak FLOPS """ | |
| # first estimate the number of flops we do per iteration. | |
| # see PaLM paper Appendix B as ref: https://arxiv.org/abs/2204.02311 | |
| N = sum(p.numel() for p in self.parameters()) | |
| cfg = self.params | |
| L, H, Q, T = cfg.n_layers, cfg.n_heads, cfg.dim//cfg.n_heads, cfg.max_seq_len | |
| flops_per_token = 6*N + 12*L*H*Q*T | |
| flops_per_fwdbwd = flops_per_token * T | |
| flops_per_iter = flops_per_fwdbwd * fwdbwd_per_iter | |
| # express our flops throughput as ratio of A100 bfloat16 peak flops | |
| flops_achieved = flops_per_iter * (1.0/dt) # per second | |
| flops_promised = 312e12 # A100 GPU bfloat16 peak flops is 312 TFLOPS | |
| mfu = flops_achieved / flops_promised | |
| return mfu | |
| def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None): | |
| """ | |
| Take a conditioning sequence of indices idx (LongTensor of shape (b,t)) and complete | |
| the sequence max_new_tokens times, feeding the predictions back into the model each time. | |
| Most likely you'll want to make sure to be in model.eval() mode of operation for this. | |
| Also note this is a super inefficient version of sampling with no key/value cache. | |
| """ | |
| for _ in range(max_new_tokens): | |
| # if the sequence context is growing too long we must crop it at block_size | |
| idx_cond = idx if idx.size(1) <= self.params.max_seq_len else idx[:, -self.params.max_seq_len:] | |
| # forward the model to get the logits for the index in the sequence | |
| logits = self(idx_cond) | |
| logits = logits[:, -1, :] # crop to just the final time step | |
| if temperature == 0.0: | |
| # "sample" the single most likely index | |
| _, idx_next = torch.topk(logits, k=1, dim=-1) | |
| else: | |
| # pluck the logits at the final step and scale by desired temperature | |
| logits = logits / temperature | |
| # optionally crop the logits to only the top k options | |
| if top_k is not None: | |
| v, _ = torch.topk(logits, min(top_k, logits.size(-1))) | |
| logits[logits < v[:, [-1]]] = -float('Inf') | |
| # apply softmax to convert logits to (normalized) probabilities | |
| probs = F.softmax(logits, dim=-1) | |
| idx_next = torch.multinomial(probs, num_samples=1) | |
| # append sampled index to the running sequence and continue | |
| idx = torch.cat((idx, idx_next), dim=1) | |
| return idx | |
| # class ConditionalTransformer(nn.Module): | |
| # last_loss: Optional[torch.Tensor] | |
| # def __init__(self, params: ConditionalModelArgs): | |
| # super().__init__() | |
| # self.params = params | |
| # self.vocab_size = params.vocab_size | |
| # self.n_layers = params.n_layers | |
| # self.cond_embeddings = ContinuousConditionalEmbedding(params) | |
| # self.tok_embeddings = nn.Embedding(params.vocab_size, params.dim) | |
| # self.dropout = nn.Dropout(params.dropout) | |
| # self.layers = torch.nn.ModuleList() | |
| # for layer_id in range(params.n_layers): | |
| # self.layers.append(TransformerBlock(layer_id, params)) | |
| # self.norm = RMSNorm(params.dim, eps=params.norm_eps) | |
| # self.output = nn.Linear(params.dim, params.vocab_size, bias=False) | |
| # # self.proj = nn.Linear(params.dim, params.dim, bias=False) | |
| # # share the unembedding parameters with the embedding parameters | |
| # self.tok_embeddings.weight = self.output.weight # https://paperswithcode.com/method/weight-tying | |
| # # some useful precompute for the RoPE relative positional embeddings | |
| # freqs_cos, freqs_sin = precompute_freqs_cis(self.params.dim // self.params.n_heads, self.params.max_seq_len) | |
| # self.register_buffer("freqs_cos", freqs_cos, persistent=False) | |
| # self.register_buffer("freqs_sin", freqs_sin, persistent=False) | |
| # # init all weights | |
| # self.apply(self._init_weights) | |
| # # apply special scaled init to the residual projections, per GPT-2 paper | |
| # for pn, p in self.named_parameters(): | |
| # if pn.endswith('w3.weight') or pn.endswith('wo.weight'): | |
| # torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * params.n_layers)) | |
| # # Initialize attribute for the loss of the last forward call. This will be set if the forward is called with a targets tensor. | |
| # self.last_loss = None | |
| # def _init_weights(self, module): | |
| # if isinstance(module, nn.Linear): | |
| # torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) | |
| # if module.bias is not None: | |
| # torch.nn.init.zeros_(module.bias) | |
| # elif isinstance(module, nn.Embedding): | |
| # torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) | |
| # def forward(self, conditions: dict, tokens: torch.Tensor, targets: Optional[torch.Tensor] = None, masks: Optional[torch.Tensor] = None) -> torch.Tensor: | |
| # tok = self.tok_embeddings(tokens) | |
| # cond = self.cond_embeddings(conditions) | |
| # h = torch.cat([cond, tok], dim=1) | |
| # _bsz, seqlen, _ = h.shape | |
| # h = self.dropout(h) | |
| # freqs_cos = self.freqs_cos[:seqlen] | |
| # freqs_sin = self.freqs_sin[:seqlen] | |
| # for layer in self.layers: | |
| # h = layer(h, freqs_cos, freqs_sin) | |
| # h = self.norm(h) | |
| # if targets is not None: | |
| # # if we are given some desired targets also calculate the loss | |
| # logits = self.output(h) | |
| # # print(logits.shape, h.shape, tokens.shape, targets.shape) | |
| # # print(tokens[0]) | |
| # # print(targets[0]) | |
| # proj = None#self.proj(h) | |
| # self.last_loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1), ignore_index=0) # ignore [PAD] | |
| # else: | |
| # # inference-time mini-optimization: only forward the output on the very last position | |
| # logits = self.output(h[:, [-1], :]) # note: using list [-1] to preserve the time dim | |
| # proj = None | |
| # self.last_loss = None | |
| # return logits | |
| # def configure_optimizers_1st_stage(self, weight_decay, learning_rate, betas, device_type): | |
| # high_param_dict = {pn: p for pn, p in self.named_parameters() if 'cond_embeddings' in pn or 'proj' in pn} | |
| # high_param_dict = {pn: p for pn, p in high_param_dict.items() if p.requires_grad} | |
| # low_param_dict = {pn: p for pn, p in self.named_parameters() if pn not in high_param_dict} | |
| # low_param_dict = {pn: p for pn, p in low_param_dict.items() if p.requires_grad} | |
| # high_decay_params = [p for n, p in high_param_dict.items() if p.dim() >= 2] | |
| # high_nodecay_params = [p for n, p in high_param_dict.items() if p.dim() < 2] | |
| # low_decay_params = [p for n, p in low_param_dict.items() if p.dim() >= 2] | |
| # low_nodecay_params = [p for n, p in low_param_dict.items() if p.dim() < 2] | |
| # optim_groups = [ | |
| # {'params': high_decay_params, 'weight_decay': weight_decay, 'lr': learning_rate}, | |
| # {'params': high_nodecay_params, 'weight_decay': 0.0, 'lr': learning_rate}, | |
| # {'params': low_decay_params, 'weight_decay': weight_decay, 'lr': learning_rate / 100}, | |
| # {'params': low_nodecay_params, 'weight_decay': 0.0, 'lr': learning_rate / 100}, | |
| # ] | |
| # # Create AdamW optimizer and use the fused version if it is available | |
| # fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters | |
| # use_fused = fused_available and device_type == 'cuda' | |
| # extra_args = dict(fused=True) if use_fused else dict() | |
| # optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas, **extra_args) | |
| # print(f"using fused AdamW: {use_fused}") | |
| # return optimizer | |
| # def configure_optimizers_2nd_stage(self, weight_decay, learning_rate, betas, device_type): | |
| # high_param_dict = {pn: p for pn, p in self.named_parameters() if 'cond_embeddings' in pn} | |
| # high_param_dict = {pn: p for pn, p in high_param_dict.items() if p.requires_grad} | |
| # mid_param_dict = {pn: p for pn, p in self.named_parameters() if 'output' in pn or ('embeddings' in pn and pn not in high_param_dict)} | |
| # mid_param_dict = {pn: p for pn, p in mid_param_dict.items() if p.requires_grad} | |
| # low_param_dict = {pn: p for pn, p in self.named_parameters() if pn not in high_param_dict and pn not in mid_param_dict} | |
| # low_param_dict = {pn: p for pn, p in low_param_dict.items() if p.requires_grad} | |
| # high_decay_params = [p for n, p in high_param_dict.items() if p.dim() >= 2] | |
| # high_nodecay_params = [p for n, p in high_param_dict.items() if p.dim() < 2] | |
| # mid_decay_params = [p for n, p in mid_param_dict.items() if p.dim() >= 2] | |
| # mid_nodecay_params = [p for n, p in mid_param_dict.items() if p.dim() < 2] | |
| # low_decay_params = [p for n, p in low_param_dict.items() if p.dim() >= 2] | |
| # low_nodecay_params = [p for n, p in low_param_dict.items() if p.dim() < 2] | |
| # optim_groups = [ | |
| # {'params': high_decay_params, 'weight_decay': weight_decay, 'lr': learning_rate}, | |
| # {'params': high_nodecay_params, 'weight_decay': 0.0, 'lr': learning_rate}, | |
| # {'params': mid_decay_params, 'weight_decay': weight_decay, 'lr': learning_rate / 10}, | |
| # {'params': mid_nodecay_params, 'weight_decay': 0.0, 'lr': learning_rate / 10}, | |
| # {'params': low_decay_params, 'weight_decay': weight_decay, 'lr': learning_rate / 100}, | |
| # {'params': low_nodecay_params, 'weight_decay': 0.0, 'lr': learning_rate / 100}, | |
| # ] | |
| # # Create AdamW optimizer and use the fused version if it is available | |
| # fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters | |
| # use_fused = fused_available and device_type == 'cuda' | |
| # extra_args = dict(fused=True) if use_fused else dict() | |
| # optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas, **extra_args) | |
| # print(f"using fused AdamW: {use_fused}") | |
| # return optimizer | |
| # def configure_optimizers_3rd_stage(self, weight_decay, learning_rate, betas, device_type): | |
| # # start with all of the candidate parameters | |
| # param_dict = {pn: p for pn, p in self.named_parameters()} | |
| # # filter out those that do not require grad | |
| # param_dict = {pn: p for pn, p in param_dict.items() if p.requires_grad} | |
| # # create optim groups. Any parameters that is 2D will be weight decayed, otherwise no. | |
| # # i.e. all weight tensors in matmuls + embeddings decay, all biases and layernorms don't. | |
| # decay_params = [p for n, p in param_dict.items() if p.dim() >= 2] | |
| # nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2] | |
| # optim_groups = [ | |
| # {'params': decay_params, 'weight_decay': weight_decay}, | |
| # {'params': nodecay_params, 'weight_decay': 0.0} | |
| # ] | |
| # # Create AdamW optimizer and use the fused version if it is available | |
| # fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters | |
| # use_fused = fused_available and device_type == 'cuda' | |
| # extra_args = dict(fused=True) if use_fused else dict() | |
| # optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas, **extra_args) | |
| # print(f"using fused AdamW: {use_fused}") | |
| # return optimizer | |
| # def estimate_mfu(self, fwdbwd_per_iter, dt): | |
| # """ estimate model flops utilization (MFU) in units of A100 bfloat16 peak FLOPS """ | |
| # # first estimate the number of flops we do per iteration. | |
| # # see PaLM paper Appendix B as ref: https://arxiv.org/abs/2204.02311 | |
| # N = sum(p.numel() for p in self.parameters()) | |
| # cfg = self.params | |
| # L, H, Q, T = cfg.n_layers, cfg.n_heads, cfg.dim//cfg.n_heads, cfg.max_seq_len | |
| # flops_per_token = 6*N + 12*L*H*Q*T | |
| # flops_per_fwdbwd = flops_per_token * T | |
| # flops_per_iter = flops_per_fwdbwd * fwdbwd_per_iter | |
| # # express our flops throughput as ratio of A100 bfloat16 peak flops | |
| # flops_achieved = flops_per_iter * (1.0/dt) # per second | |
| # flops_promised = 312e12 # A100 GPU bfloat16 peak flops is 312 TFLOPS | |
| # mfu = flops_achieved / flops_promised | |
| # return mfu | |
| # @torch.inference_mode() | |
| # def generate(self, conditions, idx, max_new_tokens, temperature=1.0, top_k=None): | |
| # """ | |
| # Take a conditioning sequence of indices idx (LongTensor of shape (b,t)) and complete | |
| # the sequence max_new_tokens times, feeding the predictions back into the model each time. | |
| # Most likely you'll want to make sure to be in model.eval() mode of operation for this. | |
| # Also note this is a super inefficient version of sampling with no key/value cache. | |
| # """ | |
| # for _ in range(max_new_tokens): | |
| # # if the sequence context is growing too long we must crop it at block_size | |
| # idx_cond = idx if idx.size(1) <= self.params.max_seq_len else idx[:, -self.params.max_seq_len:] | |
| # # forward the model to get the logits for the index in the sequence | |
| # logits = self(conditions, idx_cond) | |
| # logits = logits[:, -1, :] # crop to just the final time step | |
| # if temperature == 0.0: | |
| # # "sample" the single most likely index | |
| # _, idx_next = torch.topk(logits, k=1, dim=-1) | |
| # else: | |
| # # pluck the logits at the final step and scale by desired temperature | |
| # logits = logits / temperature | |
| # # optionally crop the logits to only the top k options | |
| # if top_k is not None: | |
| # v, _ = torch.topk(logits, min(top_k, logits.size(-1))) | |
| # logits[logits < v[:, [-1]]] = -float('Inf') | |
| # # apply softmax to convert logits to (normalized) probabilities | |
| # probs = F.softmax(logits, dim=-1) | |
| # idx_next = torch.multinomial(probs, num_samples=1) | |
| # # append sampled index to the running sequence and continue | |
| # idx = torch.cat((idx, idx_next), dim=1) | |
| # return idx |