bartholomort/nanoGPT-BitNet158b-Lua

1"""2Full definition of a GPT Language Model, all of it in this single file.3References:41) the official GPT-2 TensorFlow implementation released by OpenAI:5https://github.com/openai/gpt-2/blob/master/src/model.py62) huggingface/transformers PyTorch implementation:7https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt2/modeling_gpt2.py8"""9 10import math11import inspect12from dataclasses import dataclass13 14import torch15import torch.nn as nn16from torch.nn import functional as F17 18# https://github.com/suito555/bitnet158b19class BitLinear158b(nn.Linear):20    def __init__(self, in_features, out_features, bias=False, bit_scale = 8):21        super(BitLinear158b, self).__init__(in_features, out_features, bias)22        self.bit_scale = bit_scale23        self.Q_b = 2 ** (self.bit_scale - 1)24        self.eps = 1e-825        26    def quantize_activations(self, input_norm, abs_max_x_value):27        scaled_x = input_norm * self.Q_b / (abs_max_x_value + self.eps)28        quantized_x = torch.round(torch.clamp(29            scaled_x, -self.Q_b + self.eps, self.Q_b - self.eps30        ))31        # Inserting following comment outed line prevents the loss to be NaN if you use torch.compile with torch==2.0.0 or torch==2.2.132        # If you use torch==2.1.2, it will not be compiled.33        if(torch.isnan(scaled_x).any()): 34            print("Nan!")35        #STE36        quantized_x = (quantized_x - scaled_x).detach() + scaled_x37        return quantized_x38    39    def ternarize_weights(self,abs_mean_W_value):40        scaled_W = self.weight / (abs_mean_W_value + self.eps)41        quantize_weights = torch.clamp(scaled_W.round(), -1, 1)42        #STE 43        quantize_weights = (quantize_weights - self.weight).detach() + self.weight44        return quantize_weights45 46    def forward(self, input):47        input_norm = F.layer_norm(input, (self.in_features,))48        49        abs_max_x_value = input_norm.abs().max() #gamma50        quant_scaled_input = self.quantize_activations(input_norm,abs_max_x_value)51 52        abs_mean_W_value = self.weight.abs().mean() #beta53        ternarized_weights = self.ternarize_weights(abs_mean_W_value)54 55        matmal_weight = F.linear(quant_scaled_input, ternarized_weights, self.bias)56 57        beta_gamma = abs_mean_W_value * abs_max_x_value58        output = matmal_weight * beta_gamma / self.Q_b59        return output60 61# class BitLinear158bNIQ(nn.Linear):62#     def __init__(self, in_features, out_features, bias=False, bit_scale = 8):63#         super(BitLinear158bNIQ, self).__init__(in_features, out_features, bias)64#         self.bit_scale = bit_scale65#         self.Q_b = 2 ** (self.bit_scale - 1)66#         self.eps = 1e-867        68#     def quantize_activations(self, input_norm, abs_max_x_value):69#         scaled_x = torch.clamp(70#             input_norm * self.Q_b / (abs_max_x_value + self.eps), -self.Q_b + self.eps, self.Q_b - self.eps71#         )72#         return scaled_x73    74#     def ternarize_weights(self,abs_mean_W_value):75#         scaled_W = self.weight / (abs_mean_W_value + self.eps)76#         quantize_weights = torch.sign(torch.clamp(scaled_W.round(), -1, 1))77#         #STE 78#         quantize_weights = (quantize_weights - self.weight).detach() + self.weight79#         return quantize_weights80 81#     def forward(self, input):82#         input_norm = F.layer_norm(input, (self.in_features,))83 84#         abs_mean_W_value = self.weight.abs().mean() #beta85#         ternarized_weights = self.ternarize_weights(abs_mean_W_value)86 87#         matmal_weight = F.linear(input_norm, ternarized_weights, self.bias)88 89#         beta_gamma = abs_mean_W_value90#         output = matmal_weight * beta_gamma / self.Q_b91#         return output92 93class LayerNorm(nn.Module):94    """ LayerNorm but with an optional bias. PyTorch doesn't support simply bias=False """95 96    def __init__(self, ndim, bias):97        super().__init__()98        self.weight = nn.Parameter(torch.ones(ndim))99        self.bias = nn.Parameter(torch.zeros(ndim)) if bias else None100 101    def forward(self, input):102        return F.layer_norm(input, self.weight.shape, self.weight, self.bias, 1e-5)103 104class CausalSelfAttention(nn.Module):105 106    def __init__(self, config):107        super().__init__()108        assert config.n_embd % config.n_head == 0109        # key, query, value projections for all heads, but in a batch110        self.c_attn = BitLinear158b(config.n_embd, 3 * config.n_embd, bias=config.bias)111        # output projection112        self.c_proj = BitLinear158b(config.n_embd, config.n_embd, bias=config.bias)113        # regularization114        self.attn_dropout = nn.Dropout(config.dropout)115        self.resid_dropout = nn.Dropout(config.dropout)116        self.n_head = config.n_head117        self.n_embd = config.n_embd118        self.dropout = config.dropout119        # flash attention make GPU go brrrrr but support is only in PyTorch >= 2.0120        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention')121        if not self.flash:122            print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")123            # causal mask to ensure that attention is only applied to the left in the input sequence124            self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size))125                                        .view(1, 1, config.block_size, config.block_size))126 127    def forward(self, x):128        B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd)129 130        # calculate query, key, values for all heads in batch and move head forward to be the batch dim131        q, k, v  = self.c_attn(x).split(self.n_embd, dim=2)132        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)133        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)134        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)135 136        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)137        if self.flash:138            # efficient attention using Flash Attention CUDA kernels139            y = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=self.dropout if self.training else 0, is_causal=True)140        else:141            # manual implementation of attention142            att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))143            att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))144            att = F.softmax(att, dim=-1)145            att = self.attn_dropout(att)146            y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)147        y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side148 149        # output projection150        y = self.resid_dropout(self.c_proj(y))151        return y152 153class MLP(nn.Module):154 155    def __init__(self, config):156        super().__init__()157        self.c_fc    = BitLinear158b(config.n_embd, 4 * config.n_embd, bias=config.bias)158        self.gelu    = nn.GELU()159        self.c_proj  = BitLinear158b(4 * config.n_embd, config.n_embd, bias=config.bias)160        self.dropout = nn.Dropout(config.dropout)161 162    def forward(self, x):163        x = self.c_fc(x)164        x = self.gelu(x)165        x = self.c_proj(x)166        x = self.dropout(x)167        return x168 169class Block(nn.Module):170 171    def __init__(self, config):172        super().__init__()173        self.ln_1 = LayerNorm(config.n_embd, bias=config.bias)174        self.attn = CausalSelfAttention(config)175        self.ln_2 = LayerNorm(config.n_embd, bias=config.bias)176        self.mlp = MLP(config)177 178    def forward(self, x):179        x = x + self.attn(self.ln_1(x))180        x = x + self.mlp(self.ln_2(x))181        return x182 183@dataclass184class GPTConfig:185    block_size: int = 1024186    vocab_size: int = 50304 # GPT-2 vocab_size of 50257, padded up to nearest multiple of 64 for efficiency187    n_layer: int = 12188    n_head: int = 12189    n_embd: int = 768190    dropout: float = 0.0191    bias: bool = True # True: bias in Linears and LayerNorms, like GPT-2. False: a bit better and faster192 193class GPT(nn.Module):194 195    def __init__(self, config):196        super().__init__()197        assert config.vocab_size is not None198        assert config.block_size is not None199        self.config = config200 201        self.transformer = nn.ModuleDict(dict(202            wte = nn.Embedding(config.vocab_size, config.n_embd),203            wpe = nn.Embedding(config.block_size, config.n_embd),204            drop = nn.Dropout(config.dropout),205            h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),206            ln_f = LayerNorm(config.n_embd, bias=config.bias),207        ))208        self.lm_head = BitLinear158b(config.n_embd, config.vocab_size, bias=False)209        # with weight tying when using torch.compile() some warnings get generated:210        # "UserWarning: functional_call was passed multiple values for tied weights.211        # This behavior is deprecated and will be an error in future versions"212        # not 100% sure what this is, so far seems to be harmless. TODO investigate213        self.transformer.wte.weight = self.lm_head.weight # https://paperswithcode.com/method/weight-tying214 215        # init all weights216        self.apply(self._init_weights)217        # apply special scaled init to the residual projections, per GPT-2 paper218        for pn, p in self.named_parameters():219            if pn.endswith('c_proj.weight'):220                torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * config.n_layer))221 222        # report number of parameters223        print("number of parameters: %.2fM" % (self.get_num_params()/1e6,))224 225    def get_num_params(self, non_embedding=True):226        """227        Return the number of parameters in the model.228        For non-embedding count (default), the position embeddings get subtracted.229        The token embeddings would too, except due to the parameter sharing these230        params are actually used as weights in the final layer, so we include them.231        """232        n_params = sum(p.numel() for p in self.parameters())233        if non_embedding:234            n_params -= self.transformer.wpe.weight.numel()235        return n_params236 237    def _init_weights(self, module):238        if isinstance(module, nn.Linear):239            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)240            if module.bias is not None:241                torch.nn.init.zeros_(module.bias)242        elif isinstance(module, BitLinear158b):243            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)244            if module.bias is not None:245                torch.nn.init.zeros_(module.bias)246        elif isinstance(module, nn.Embedding):247            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)248 249    def forward(self, idx, targets=None):250        device = idx.device251        b, t = idx.size()252        assert t <= self.config.block_size, f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"253        pos = torch.arange(0, t, dtype=torch.long, device=device) # shape (t)254 255        # forward the GPT model itself256        tok_emb = self.transformer.wte(idx) # token embeddings of shape (b, t, n_embd)257        pos_emb = self.transformer.wpe(pos) # position embeddings of shape (t, n_embd)258        x = self.transformer.drop(tok_emb + pos_emb)259        for block in self.transformer.h:260            x = block(x)261        x = self.transformer.ln_f(x)262 263        if targets is not None:264            # if we are given some desired targets also calculate the loss265            logits = self.lm_head(x)266            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)267        else:268            # inference-time mini-optimization: only forward the lm_head on the very last position269            logits = self.lm_head(x[:, [-1], :]) # note: using list [-1] to preserve the time dim270            loss = None271 272        return logits, loss273 274    def crop_block_size(self, block_size):275        # model surgery to decrease the block size if necessary276        # e.g. we may load the GPT2 pretrained model checkpoint (block size 1024)277        # but want to use a smaller block size for some smaller, simpler model278        assert block_size <= self.config.block_size279        self.config.block_size = block_size280        self.transformer.wpe.weight = nn.Parameter(self.transformer.wpe.weight[:block_size])281        for block in self.transformer.h:282            if hasattr(block.attn, 'bias'):283                block.attn.bias = block.attn.bias[:,:,:block_size,:block_size]284 285    @classmethod286    def from_pretrained(cls, model_type, override_args=None):287        assert model_type in {'gpt2', 'gpt2-medium', 'gpt2-large', 'gpt2-xl'}288        override_args = override_args or {} # default to empty dict289        # only dropout can be overridden see more notes below290        assert all(k == 'dropout' for k in override_args)291        from transformers import GPT2LMHeadModel292        print("loading weights from pretrained gpt: %s" % model_type)293 294        # n_layer, n_head and n_embd are determined from model_type295        config_args = {296            'gpt2':         dict(n_layer=12, n_head=12, n_embd=768),  # 124M params297            'gpt2-medium':  dict(n_layer=24, n_head=16, n_embd=1024), # 350M params298            'gpt2-large':   dict(n_layer=36, n_head=20, n_embd=1280), # 774M params299            'gpt2-xl':      dict(n_layer=48, n_head=25, n_embd=1600), # 1558M params300        }[model_type]301        print("forcing vocab_size=50257, block_size=1024, bias=True")302        config_args['vocab_size'] = 50257 # always 50257 for GPT model checkpoints303        config_args['block_size'] = 1024 # always 1024 for GPT model checkpoints304        config_args['bias'] = True # always True for GPT model checkpoints305        # we can override the dropout rate, if desired306        if 'dropout' in override_args:307            print(f"overriding dropout rate to {override_args['dropout']}")308            config_args['dropout'] = override_args['dropout']309        # create a from-scratch initialized minGPT model310        config = GPTConfig(**config_args)311        model = GPT(config)312        sd = model.state_dict()313        sd_keys = sd.keys()314        sd_keys = [k for k in sd_keys if not k.endswith('.attn.bias')] # discard this mask / buffer, not a param315 316        # init a huggingface/transformers model317        model_hf = GPT2LMHeadModel.from_pretrained(model_type)318        sd_hf = model_hf.state_dict()319 320        # copy while ensuring all of the parameters are aligned and match in names and shapes321        sd_keys_hf = sd_hf.keys()322        sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.masked_bias')] # ignore these, just a buffer323        sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.bias')] # same, just the mask (buffer)324        transposed = ['attn.c_attn.weight', 'attn.c_proj.weight', 'mlp.c_fc.weight', 'mlp.c_proj.weight']325        # basically the openai checkpoints use a "Conv1D" module, but we only want to use a vanilla Linear326        # this means that we have to transpose these weights when we import them327        assert len(sd_keys_hf) == len(sd_keys), f"mismatched keys: {len(sd_keys_hf)} != {len(sd_keys)}"328        for k in sd_keys_hf:329            if any(k.endswith(w) for w in transposed):330                # special treatment for the Conv1D weights we need to transpose331                assert sd_hf[k].shape[::-1] == sd[k].shape332                with torch.no_grad():333                    sd[k].copy_(sd_hf[k].t())334            else:335                # vanilla copy over the other parameters336                assert sd_hf[k].shape == sd[k].shape337                with torch.no_grad():338                    sd[k].copy_(sd_hf[k])339 340        return model341 342    def configure_optimizers(self, weight_decay, learning_rate, betas, device_type):343        # start with all of the candidate parameters344        param_dict = {pn: p for pn, p in self.named_parameters()}345        # filter out those that do not require grad346        param_dict = {pn: p for pn, p in param_dict.items() if p.requires_grad}347        # create optim groups. Any parameters that is 2D will be weight decayed, otherwise no.348        # i.e. all weight tensors in matmuls + embeddings decay, all biases and layernorms don't.349        decay_params = [p for n, p in param_dict.items() if p.dim() >= 2]350        nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2]351        optim_groups = [352            {'params': decay_params, 'weight_decay': weight_decay},353            {'params': nodecay_params, 'weight_decay': 0.0}354        ]355        num_decay_params = sum(p.numel() for p in decay_params)356        num_nodecay_params = sum(p.numel() for p in nodecay_params)357        print(f"num decayed parameter tensors: {len(decay_params)}, with {num_decay_params:,} parameters")358        print(f"num non-decayed parameter tensors: {len(nodecay_params)}, with {num_nodecay_params:,} parameters")359        # Create AdamW optimizer and use the fused version if it is available360        fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters361        use_fused = fused_available and device_type == 'cuda'362        extra_args = dict(fused=True) if use_fused else dict()363        optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas, **extra_args)364        print(f"using fused AdamW: {use_fused}")365 366        return optimizer367 368    def estimate_mfu(self, fwdbwd_per_iter, dt):369        """ estimate model flops utilization (MFU) in units of A100 bfloat16 peak FLOPS """370        # first estimate the number of flops we do per iteration.371        # see PaLM paper Appendix B as ref: https://arxiv.org/abs/2204.02311372        N = self.get_num_params()373        cfg = self.config374        L, H, Q, T = cfg.n_layer, cfg.n_head, cfg.n_embd//cfg.n_head, cfg.block_size375        flops_per_token = 6*N + 12*L*H*Q*T376        flops_per_fwdbwd = flops_per_token * T377        flops_per_iter = flops_per_fwdbwd * fwdbwd_per_iter378        # express our flops throughput as ratio of A100 bfloat16 peak flops379        flops_achieved = flops_per_iter * (1.0/dt) # per second380        flops_promised = 312e12 # A100 GPU bfloat16 peak flops is 312 TFLOPS381        mfu = flops_achieved / flops_promised382        return mfu383 384    @torch.no_grad()385    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):386        """387        Take a conditioning sequence of indices idx (LongTensor of shape (b,t)) and complete388        the sequence max_new_tokens times, feeding the predictions back into the model each time.389        Most likely you'll want to make sure to be in model.eval() mode of operation for this.390        """391        for _ in range(max_new_tokens):392            # if the sequence context is growing too long we must crop it at block_size393            idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]394            # forward the model to get the logits for the index in the sequence395            logits, _ = self(idx_cond)396            # pluck the logits at the final step and scale by desired temperature397            logits = logits[:, -1, :] / temperature398            # optionally crop the logits to only the top k options399            if top_k is not None:400                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))401                logits[logits < v[:, [-1]]] = -float('Inf')402            # apply softmax to convert logits to (normalized) probabilities403            probs = F.softmax(logits, dim=-1)404            # sample from the distribution405            idx_next = torch.multinomial(probs, num_samples=1)406            # append sampled index to the running sequence and continue407            idx = torch.cat((idx, idx_next), dim=1)408 409        return idx410