Update README.md

README.md

@@ -50,10 +50,29 @@ This model has been 4-bit quantized Llada-8B-Base model with [GPTQModel](https:/
 
 ## Example:
 ```python
+
+# Copyright 2024-2025 ModelCloud.ai
+# Copyright 2024-2025 [email protected]
+# Contact: [email protected], x.com/qubitium
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import torch
 from datasets import load_dataset
-from gptqmodel import GPTQModel, QuantizeConfig,  BACKEND
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from gptqmodel import  GPTQModel, QuantizeConfig,  BACKEND
+from gptqmodel.models.base import BaseGPTQModel 
+from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModel
+from gptqmodel.models.auto import MODEL_MAP 
 import torch.nn.functional as F 
 import numpy as np
 
@@ -61,6 +80,64 @@ import numpy as np
 
 
 
+pretrained_model_id = '/home/chentianqi/model/GSAI-ML/LLaDA-8B-Instruct' # "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+quantized_model_id = "FunAGI/LLaDA-8B-Instruct-gptqmodel-4bit"
+
+
+
+class LladaGPTQ(BaseGPTQModel):
+    # Non-repeating layers at the root level: same level as `layers_node`
+    # Excluding `layers_node`.
+    base_modules = ["model.transformer.wte", "model.transformer.ln_f"]
+    pre_lm_head_norm_module = "model.transformer.ln_f"
+    lm_head = "model.transformer.ff_out"
+    # Below describes all the repeating layers in this transformer model
+    # `model.layers` is a node/module that hold all the repeating layers. The parent node for all n-layers.
+    layers_node = "model.transformer.blocks"
+    # Each repeating layer in `model.layers` is of type `LlamaDecoderLayer`
+    layer_type = "LLaDALlamaBlock"
+    # Inside each `LlamaDecoderLayer` layer are many internal modules
+    # List them in the order executed in model forward() code
+    # Many models have same execution order of: attention (q_k_v) projection, attention (output) projection, mlp (n) projections
+    layer_modules = [
+        ["attn_out", "k_proj", "v_proj", "q_proj"],
+        ["ff_proj", "up_proj"],
+        ["ff_out"],
+    ]
+MODEL_MAP ["llada"] = LladaGPTQ
+
+# os.makedirs(quantized_model_dir, exist_ok=True)
+def get_wikitext2(tokenizer, nsamples, seqlen):
+    traindata = load_dataset("wikitext", "wikitext-2-raw-v1", split="train").filter(
+        lambda x: len(x["text"]) >= seqlen)
+
+    return [tokenizer(example["text"]) for example in traindata.select(range(nsamples))]
+
+
+@torch.no_grad()
+def calculate_avg_ppl(model, tokenizer):
+    from gptqmodel.utils import Perplexity
+
+    ppl = Perplexity(
+        model=model,
+        tokenizer=tokenizer,
+        dataset_path="wikitext",
+        dataset_name="wikitext-2-raw-v1",
+        split="train",
+        text_column="text",
+    )
+
+    all = ppl.calculate(n_ctx=512, n_batch=512)
+
+    # average ppl
+    avg = sum(all) / len(all)
+
+    return avg
+
+dynamic = { 
+     
+ } 
+
 def add_gumbel_noise(logits, temperature):
     '''
     The Gumbel max is a method for sampling categorical distributions.
@@ -93,6 +170,82 @@ def get_num_transfer_tokens(mask_index, steps):
 
     return num_transfer_tokens
 
+def forward_process(batch, prompt_index, mask_id):
+    b, l = batch.shape
+
+    target_len = (l - prompt_index.sum()).item()
+    k = torch.randint(1, target_len + 1, (), device=batch.device)
+
+    x = torch.round(torch.linspace(float(k), k + (b - 1) * (target_len / b), steps=b, device=batch.device)).long()
+    x = ((x - 1) % target_len) + 1
+    assert x.min() >= 1 and x.max() <= target_len
+
+    indices = torch.arange(target_len, device=batch.device).repeat(b, 1)
+    is_mask = indices < x.unsqueeze(1)
+    for i in range(b):
+        is_mask[i] = is_mask[i][torch.randperm(target_len)]
+
+    is_mask = torch.cat((torch.zeros(b, prompt_index.sum(), dtype=torch.bool, device=batch.device), is_mask), dim=1)
+    noisy_batch = torch.where(is_mask, mask_id, batch)
+
+    # Return the masked batch and the mask ratio
+    return noisy_batch, (x / target_len).unsqueeze(1).repeat(1, l)
+
+
+def get_logits(model, batch, prompt_index, cfg_scale, mask_id):
+    if cfg_scale > 0.:
+        assert len(prompt_index) == batch.shape[1]
+        prompt_index = prompt_index.unsqueeze(0).repeat(batch.shape[0], 1)
+        un_batch = batch.clone()
+        un_batch[prompt_index] = mask_id
+        batch = torch.cat([batch, un_batch])
+
+    input = batch
+    logits = model(input).logits
+
+    if cfg_scale > 0.:
+        logits, un_logits = torch.chunk(logits, 2, dim=0)
+        logits = un_logits + (cfg_scale + 1) * (logits - un_logits)
+    return logits
+
+
+
+@ torch.no_grad()
+def get_log_likelihood(model, prompt, answer, mc_num=128, batch_size=32, cfg_scale=0., mask_id=126336):
+    '''
+    Args:
+        model: Mask predictor.
+        prompt: A tensor of shape (l1).
+        answer: A tensor of shape (l2).
+        mc_num: Monte Carlo estimation times.
+                As detailed in Appendix B.5. Since MMLU, CMMLU, and C-EVAL only require the likelihood of a single token, a
+                single Monte Carlo estimate is sufficient for these benchmarks. For all other benchmarks, we find that 128
+                Monte Carlo samples are adequate to produce stable results.
+        batch_size: Mini batch size.
+        cfg_scale: Unsupervised classifier-free guidance scale.
+        mask_id: The toke id of [MASK] is 126336.
+    '''
+
+    seq = torch.concatenate([prompt, answer])[None, :]
+    seq = seq.repeat((batch_size, 1)).to(model.device)
+    prompt_index = torch.arange(seq.shape[1], device=model.device) < len(prompt)
+
+    loss_ = []
+    for _ in range(mc_num // batch_size):
+  
+        perturbed_seq, p_mask = forward_process(seq, prompt_index, mask_id)
+        mask_index = perturbed_seq == mask_id
+
+        logits = get_logits(model, perturbed_seq, prompt_index, cfg_scale, mask_id)
+
+        loss = F.cross_entropy(logits[mask_index], seq[mask_index], reduction='none') / p_mask[mask_index]
+        loss = loss.sum() / batch_size
+
+        loss_.append(loss.item())
+
+    return - sum(loss_) / len(loss_)
+
+
 
 
 
@@ -167,29 +320,43 @@ def generate(model, prompt, steps=128, gen_length=128, block_length=128, tempera
     return x
 
 def main():
-    quantized_model_id="FunAGI/LLaDA-8B-Instruct-gptqmodel-4bit"
-    tokenizer = AutoTokenizer.from_pretrained(quantized_model_id ,use_fast=False)
-    
+    tokenizer = AutoTokenizer.from_pretrained(pretrained_model_id, use_fast=False)
+
+    traindataset = get_wikitext2(tokenizer, nsamples=128, seqlen=1024)
 
+    quantize_config = QuantizeConfig(
+        dynamic=dynamic,
+        bits=8,  # quantize model to 4-bit
+        group_size=128,  # it is recommended to set the value to 128,
+        desc_act = True, 
+        sym=False
+    )
     device = "cuda:0" if torch.cuda.is_available() else "cpu"
-    prompt = "Paul is at a train station and is waiting for his train. He isn't sure how long he needs to wait, but he knows that the fourth train scheduled to arrive at the station is the one he needs to get on. The first train is scheduled to arrive in 10 minutes, and this train will stay in the station for 20 minutes. The second train is to arrive half an hour after the first train leaves the station, and this second train will stay in the station for a quarter of the amount of time that the first train stayed in the station. The third train is to arrive an hour after the second train leaves the station, and this third train is to leave the station immediately after it arrives. The fourth train will arrive 20 minutes after the third train leaves, and this is the train Paul will board. In total, how long, in minutes, will Paul wait for his train?"
+    prompt = "Lily can run 12 kilometers per hour for 4 hours. After that, she runs 6 kilometers per hour. How many kilometers can she run in 8 hours?"
 
-    # # # Add special tokens for the Instruct model. The Base model does not require the following two lines.
+    # # Add special tokens for the Instruct model. The Base model does not require the following two lines.
     m = [{"role": "user", "content": prompt}, ]
     prompt = tokenizer.apply_chat_template(m, add_generation_prompt=True, tokenize=False)
 
     input_ids = tokenizer(prompt)['input_ids']
     input_ids = torch.tensor(input_ids).to(device).unsqueeze(0)
 
-    
 
 
+    device = "cuda:0" if torch.cuda.is_available() else "cpu"
     model = GPTQModel.load(quantized_model_id, device=device , trust_remote_code=True    )
 
+    steps=128
+    out = generate(model, input_ids, steps=steps , gen_length=128, block_length=32, temperature=0., cfg_scale=0., remasking='low_confidence')
+    print("*"*30+ f"GPTQ-4bit Steps {steps}"+ "*"*30)
+    print(input_ids.shape)
+    print( tokenizer.batch_decode(out[:, input_ids.shape[1]:], skip_special_tokens=True)[0])
+    del model 
+
+    model =AutoModel.from_pretrained(pretrained_model_id, trust_remote_code=True     ).cuda()
 
-    steps=256
-    out = generate(model, input_ids, steps=steps , gen_length=256, block_length=8, temperature=0., cfg_scale=0., remasking='low_confidence')
-    print("*"*30+ f"Steps {steps}"+ "*"*30)
+    out = generate(model, input_ids, steps=steps , gen_length=128, block_length=32, temperature=0., cfg_scale=0., remasking='low_confidence')
+    print("*"*30+ f"FP16 Steps {steps}"+ "*"*30)
     print(input_ids.shape)
     print( tokenizer.batch_decode(out[:, input_ids.shape[1]:], skip_special_tokens=True)[0])
 
@@ -206,4 +373,6 @@ if __name__ == "__main__":
 
     main()
 
+
+
 ```