nvidia/Efficient-DLM-8B

1import numpy as np2import torch3import torch.nn.functional as F4 5 6def add_gumbel_noise(logits, temperature):7    '''8    The Gumbel max is a method for sampling categorical distributions.9    According to arXiv:2409.02908, for MDM, low-precision Gumbel Max improves perplexity score but reduces generation quality.10    Thus, we use float64.11    '''12    if temperature == 0:13        return logits14    logits = logits.to(torch.float64)15    noise = torch.rand_like(logits, dtype=torch.float64)16    gumbel_noise = (- torch.log(noise)) ** temperature17    return logits.exp() / gumbel_noise18 19 20def get_transfer_index(logits, temperature, remasking, mask_index, x, num_transfer_tokens, threshold=None,neg_entropy=False):21    logits_with_noise = add_gumbel_noise(logits, temperature=temperature)22    x0 = torch.argmax(logits_with_noise, dim=-1)23 24    if remasking == 'low_confidence':25        # p = F.softmax(logits.to(torch.float64), dim=-1)26        p = F.softmax(logits, dim=-1)27        x0_p = torch.squeeze(28            torch.gather(p, dim=-1, index=torch.unsqueeze(x0, -1)), -1) # b, l29    elif remasking == 'top_p_margin':30        # Compute probabilities31        p = F.softmax(logits, dim=-1)                       # (B, L, V)32        # Top-2 per position33        top2 = torch.topk(p, k=2, dim=-1).values            # (B, L, 2)34        margin = top2[..., 0] - top2[..., 1]                # (B, L)35 36        # Normalize margin to [0,1] over MASKED positions per row37        plus_inf  = torch.full_like(margin, float('inf'))38        minus_inf = torch.full_like(margin, float('-inf'))39        masked_for_min = torch.where(mask_index, margin, plus_inf)40        masked_for_max = torch.where(mask_index, margin, minus_inf)41        row_min = masked_for_min.amin(dim=1, keepdim=True)  # (B, 1)42        row_max = masked_for_max.amax(dim=1, keepdim=True)  # (B, 1)43        denom = (row_max - row_min)44 45        # If denom==0 (all equal), set normalized=1 on masked; 0 elsewhere by default46        normalized = torch.zeros_like(margin)47        nonzero = denom > 048        normalized = torch.where(49            mask_index & nonzero,50            (margin - row_min) / (denom + 1e-12),51            normalized52        )53        normalized = torch.where(54            mask_index & (~nonzero),55            torch.ones_like(normalized),56            normalized57        )58        x0_p = normalized  # ∈ [0,1] on masked positions59    elif remasking == 'random':60        x0_p = torch.rand((x0.shape[0], x0.shape[1]), device=x0.device)61    else:62        raise NotImplementedError(remasking)63    64    # Calculate negative entropy if requested65    if neg_entropy:66        # p = F.softmax(logits.to(torch.float64), dim=-1)67        p = F.softmax(logits, dim=-1)68        epsilon = 1e-1069        log_probs = torch.log(p + epsilon)70        confidence_scores = torch.sum(p * log_probs, dim=-1)  # negative entropy per position71    else:72        confidence_scores = x0_p73    74    x0 = torch.where(mask_index, x0, x)75    confidence = torch.where(mask_index, confidence_scores, -np.inf)76 77    transfer_index = torch.zeros_like(x0, dtype=torch.bool, device=x0.device)78    if threshold is not None:79        num_transfer_tokens = mask_index.sum(dim=1, keepdim=True)80    # print(f'confidence: {confidence}')81    for j in range(confidence.shape[0]):82        _, select_index = torch.topk(confidence[j], k=num_transfer_tokens[j])83        transfer_index[j, select_index] = True84        if threshold is not None:85            for k in range(1, num_transfer_tokens[j]):86                if confidence[j, select_index[k]] < threshold:87                    transfer_index[j, select_index[k]] = False88    return x0, transfer_index89 90 91def get_num_transfer_tokens(mask_index, steps: int):92    mask_num = mask_index.sum(dim=1, keepdim=True)93    base = mask_num // steps94    remainder = mask_num % steps95    num_transfer_tokens = torch.zeros(mask_num.size(0), steps, device=mask_index.device, dtype=torch.int64) + base96    for i in range(mask_num.size(0)):97        num_transfer_tokens[i, : int(remainder[i])] += 198    return num_transfer_tokens99 100 101@torch.no_grad()102def generate_with_prefix_cache_block_diff(103    model,104    prompt,105    steps=128,106    gen_length=128,107    block_length=128,108    temperature=0.,109    remasking='low_confidence',110    mask_id=126336,111    threshold=None,112    factor=None,113    shift_logits=False,114    neg_entropy=False,115):116    dream_style=shift_logits117    # Initialize the accumulator118    x_accum = prompt.clone()119 120    assert gen_length % block_length == 0121    num_blocks = gen_length // block_length122 123    assert steps % num_blocks == 0124    steps_per_block = steps // num_blocks125 126    nfe = 0127 128    # Compute KV cache for the prompt initially129    output = model(prompt, use_cache=True)130    past_key_values = output.past_key_values131 132    # For dream_style: store the "next token logit" of the context133    next_logits_context = None134    if dream_style:135        next_logits_context = output.logits[:, -1:, :]  # (B, 1, V)136 137    for num_block in range(num_blocks):138        # Create a new block with mask tokens (no seeding)139        mask_block = torch.ones(140            (prompt.shape[0], block_length),141            dtype=prompt.dtype,142            device=prompt.device143        ) * mask_id144 145        # Append the block of masks146        x_accum = torch.cat([x_accum, mask_block], dim=1)147        current_block_start = prompt.size(1) + num_block * block_length148        block_slice = slice(current_block_start, current_block_start + block_length)149 150        # Build the initial mask for this block151        mask_block_idx0 = (x_accum[:, block_slice] == mask_id)  # (B, Lb)152 153        # Precompute the transfer schedule for this block154        if dream_style:155            # still denoise *all* positions (0..Lb-1), since none are seeded156            schedule_mask = mask_block_idx0157        else:158            schedule_mask = mask_block_idx0159 160        num_transfer_tokens = get_num_transfer_tokens(schedule_mask, steps_per_block)  # (B, steps)161 162        # Denoise the current block163        for i in range(steps_per_block):164            mask_block_idx = (x_accum[:, block_slice] == mask_id)  # (B, Lb)165            if mask_block_idx.sum() == 0:166                break167 168            nfe += 1169 170            # Forward only the current noisy block using cached context171            logits_block = model(172                x_accum[:, block_slice],173                past_key_values=past_key_values,174                use_cache=False175            ).logits176 177            if dream_style:178                # Align logits so that each masked position has a predictor:179                # prepend context-next logit, then use logits_block[:-1]180                if block_length == 1:181                    logits_use = next_logits_context              # (B, 1, V)182                else:183                    logits_use = torch.cat(184                        [next_logits_context, logits_block[:, :-1, :]],185                        dim=1186                    )  # (B, Lb, V)187 188                mask_use = mask_block_idx                        # (B, Lb)189                x_use   = x_accum[:, block_slice]                # (B, Lb)190 191                x0, transfer_idx = get_transfer_index(192                    logits_use, temperature, remasking, mask_use, x_use,193                    num_transfer_tokens=num_transfer_tokens[:, i],194                    threshold=threshold, neg_entropy=neg_entropy195                )196                cur = x_accum[:, block_slice].clone()197                cur[transfer_idx] = x0[transfer_idx]198                x_accum[:, block_slice] = cur199 200            else:201                # non-AR (same-position) case202                x0, transfer_idx = get_transfer_index(203                    logits_block, temperature, remasking, mask_block_idx,204                    x_accum[:, block_slice],205                    num_transfer_tokens=num_transfer_tokens[:, i],206                    threshold=threshold, neg_entropy=neg_entropy207                )208                cur = x_accum[:, block_slice].clone()209                cur[transfer_idx] = x0[transfer_idx]210                x_accum[:, block_slice] = cur211 212        # after block is fully denoised, update KV cache213        output = model(214            x_accum[:, block_slice],215            past_key_values=past_key_values,216            use_cache=True217        )218        past_key_values = output.past_key_values219        nfe += 1220 221        if dream_style and num_block < num_blocks - 1:222            # refresh context-next logit for the next block223            next_logits_context = output.logits[:, -1:, :]  # (B, 1, V)224 225    return x_accum, nfe226