Skip to content

Questions about "forward_process_semi_ar" function #7

Description

@BHmingyang

Thanks for releasing dParallel! Insightful work. I have some questions about the following fuction:
def forward_process_semi_ar(input_ids, prompt_lengths, mask_token_id=126336, block_size=32, eps=1e-3):
"""Semi-autoregressive forward masking process"""
b, l = input_ids.shape
device = input_ids.device

noisy_batch = input_ids.clone()
noisy_batch_rev = input_ids.clone()
masked_indices = torch.zeros_like(input_ids, dtype=torch.bool)
masked_indices_rev = torch.zeros_like(input_ids, dtype=torch.bool)

# prompt mask
token_positions = torch.arange(l, device=device).expand(b, l)
prompt_mask = token_positions < prompt_lengths.unsqueeze(1)


noisy_batch[prompt_mask] = input_ids[prompt_mask]
noisy_batch_rev[prompt_mask] = input_ids[prompt_mask]

# semi-autoregressive mask
for i in range(b):
    prompt_len = prompt_lengths[i].item()
    response_len = l - prompt_len
    
    if response_len > 0:

        max_blocks = response_len // block_size
        
        num_blocks_to_mask = random.randint(0, max_blocks)
        num_tokens_to_mask = num_blocks_to_mask * block_size
        
        mask_start = prompt_len + num_tokens_to_mask
        if num_blocks_to_mask == max_blocks:
            mask_end = l
        else:
            mask_end = mask_start + block_size
        
        **# t = torch.rand(b, device=input_ids.device)
        t = torch.full((b,), 0.5, device=input_ids.device)  # 50% mask ratio
        p_mask = (1 - eps) * t + eps
        seg_len = mask_end - mask_start
        p_mask = p_mask[:, None].repeat(1, seg_len)
        seg_mask = torch.rand((b, seg_len), device=input_ids.device) < p_mask
        masked_indices[:, mask_start:mask_end] = seg_mask
        masked_indices_rev[:, mask_start:mask_end] = ~seg_mask**
        
        noisy_batch = torch.where(masked_indices, 126336, input_ids)
        noisy_batch[i, mask_end:l] = mask_token_id

        noisy_batch_rev = torch.where(masked_indices_rev, 126336, input_ids)
        noisy_batch_rev[i, mask_end:l] = mask_token_id


return noisy_batch, noisy_batch_rev, masked_indices, masked_indices_rev

The loop is for i in range(b), processing each sample.nHowever, inside the loop, seg_mask = torch.rand((b, seg_len), ...) generates a mask for the entire batch, and then writes it to:
masked_indices[:, mask_start:mask_end] = seg_mask
masked_indices_rev[:, mask_start:mask_end] = ~seg_mask
This overwrites the entire batch's mask region in each iteration, ultimately only retaining the mask generated in the last iteration. This may not be a big problem, but feeling weired. Best wishes.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions