File size: 7,634 Bytes
032e687 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 | from typing import Dict, Sequence
import numpy as np
import torch
from torch.nn.utils.rnn import pad_sequence
from xtuner.parallel.sequence import (get_sequence_parallel_world_size,
pad_for_sequence_parallel)
from xtuner.utils import DEFAULT_PAD_TOKEN_INDEX, IGNORE_INDEX
from einops import rearrange
NON_VISION_TOKEN = -1
def generate_mm_pos_ids_singleit(input_ids, vpatch_id, h, w):
if h * w == 0:
nt = len(input_ids)
# pure text
position_id = torch.arange(nt).unsqueeze(-1).repeat(1, 3)
assert len(input_ids) == position_id.size(0)
position_id = rearrange(position_id, "slen d -> d slen").long()
return position_id
input_ids_pt = torch.Tensor(input_ids).int()
vpatch_pos = torch.argwhere(input_ids_pt == vpatch_id)
vpatch_start_pos = vpatch_pos[0].item()
nt = len(input_ids) - (h * w) + 1
# v_pos
t_indices = torch.arange(1)
h_indices = torch.arange(h)
w_indices = torch.arange(w)
v_pos_id = torch.stack(torch.meshgrid(t_indices, h_indices, w_indices, indexing='ij'), dim=0)
v_pos_id = rearrange(v_pos_id, "d t h w -> (t h w) d") # [h*w, 3]
v_pos_id += vpatch_start_pos
position_id = torch.cat(
[
torch.arange(vpatch_start_pos).unsqueeze(-1).repeat(1, 3),
v_pos_id,
torch.arange(nt - vpatch_start_pos - 1).unsqueeze(-1).repeat(1, 3) + v_pos_id.max() + 1,
],
dim=0
)
assert len(input_ids) == position_id.size(0)
position_id = rearrange(position_id, "slen d -> d slen").long()
return position_id
def st_collate_fn(instances: Sequence[Dict],
pad_index: int = DEFAULT_PAD_TOKEN_INDEX,
return_hf_format: bool = False,
use_varlen_attn: bool = False):
seq_parallel_world_size = get_sequence_parallel_world_size()
vision_patch_idx = instances[0].get('vision_patch_idx')
input_ids, labels = [], []
has_image = any(inst.get('vision_patches') is not None for inst in instances)
has_mask = any(inst.get('masks') is not None for inst in instances)
if use_varlen_attn:
position_ids, cumulative_len = [], []
assert len(instances) == 1, (
f'If utilizing varlen attention, the batch size should be'
f' set to 1, but got {len(instances)}')
assert not has_image, 'Currently, it is not configured to '
'accommodate the use of varlen Attention in multimodal training'
patch_nums_per_images = []
vision_start_end = []
vision_patch_indices = []
if has_image:
vision_patches = []
else:
vision_patches = None
if has_mask:
masks = []
else:
masks = None
_vision_indexes_prefix = 0
for example in instances:
input_ids.append(torch.LongTensor(example['input_ids']))
labels.append(torch.LongTensor(example['labels']))
patch_nums_per_images.append(example['patch_nums_per_images'])
vision_start_end.append(example['vision_start_end'])
# compute new multi-batch vision patch indices
batch_vision_patch_indices = torch.LongTensor(example['vision_patch_indices'])
batch_vision_patch_indices[batch_vision_patch_indices!=NON_VISION_TOKEN] += _vision_indexes_prefix
_vision_indexes_prefix = max(torch.max(batch_vision_patch_indices), 0)
vision_patch_indices.append(batch_vision_patch_indices)
if use_varlen_attn:
cumulative_len.append(torch.IntTensor(example['cumulative_len']))
position_ids.append(torch.LongTensor(example['position_ids']))
if has_image:
if 'vision_patches' in example.keys():
vision_patches.append(example['vision_patches'])
if has_mask:
if 'masks' in example.keys() and example['masks'] is not None:
masks.append(example['masks'])
else:
masks.append(None)
ori_length = [len(ids) for ids in input_ids]
if len(instances) > 1:
input_ids = pad_sequence(
input_ids, batch_first=True, padding_value=pad_index)
labels = pad_sequence(
labels, batch_first=True, padding_value=IGNORE_INDEX)
vision_patch_indices = pad_sequence(
vision_patch_indices, batch_first=True, padding_value=NON_VISION_TOKEN)
else:
input_ids = torch.stack(input_ids)
labels = torch.stack(labels)
vision_patch_indices = torch.stack(vision_patch_indices)
if use_varlen_attn:
assert input_ids.size(1) % seq_parallel_world_size == 0
attention_mask = None
position_ids = torch.stack(position_ids, dim=0)
else:
# Some tokenizers have the same eos token and pad token, so input_ids
# cannot be masked directly based on the pad token id.
attention_mask = torch.zeros(input_ids.shape[0], 1, input_ids.shape[1], input_ids.shape[1]).bool()
for i, length in enumerate(ori_length):
attention_mask[i, 0, :length, :length] = create_single_prefix_mask(vision_start_end[i], length)
bs, seq_len = input_ids.shape
position_ids = []
for input_id, patch_nums_per_image in zip(input_ids, patch_nums_per_images):
position_id = generate_mm_pos_ids_singleit(
input_id.cpu().numpy().tolist(), vision_patch_idx,
patch_nums_per_image[0], patch_nums_per_image[1])
position_ids.append(position_id)
position_ids = torch.stack(position_ids, dim=1)
if seq_parallel_world_size > 1:
input_ids = pad_for_sequence_parallel(input_ids, pad_index)
labels = pad_for_sequence_parallel(labels, IGNORE_INDEX)
position_ids = pad_for_sequence_parallel(position_ids, 0)
if attention_mask is not None:
attention_mask = pad_for_sequence_parallel(attention_mask, 0)
if has_image:
if len(vision_patches) == 0:
vision_patches = None
else:
vision_patches = torch.cat(vision_patches, dim=0)
if use_varlen_attn:
max_seqlen = (
cumulative_len[0][1:] - # noqa: W504
cumulative_len[0][:-1]).max().item()
data_dict = {
'input_ids': input_ids,
'cumulative_len': cumulative_len,
'position_ids': position_ids,
'labels': labels,
'max_seqlen': max_seqlen,
'vision_patch_indices': vision_patch_indices,
'masks': masks,
'vision_patches': vision_patches,
'patch_nums_per_images': patch_nums_per_images
}
else:
data_dict = {
'input_ids': input_ids,
'attention_mask': attention_mask,
'position_ids': position_ids,
'labels': labels,
'vision_patch_indices': vision_patch_indices,
'masks': masks,
'vision_patches': vision_patches,
'patch_nums_per_images': patch_nums_per_images
}
if return_hf_format:
return data_dict
else:
return {'data': data_dict, 'data_samples': None}
def create_single_prefix_mask(vision_start_end, max_len):
if vision_start_end is None:
# pure text
attn_mask = torch.tril(torch.ones(max_len, max_len))
else:
attn_mask = torch.zeros(max_len, max_len)
attn_mask[vision_start_end[0]-1:vision_start_end[1]+1, vision_start_end[0]-1:vision_start_end[1]+1] = 1
causal_mask = torch.tril(torch.ones(max_len, max_len))
attn_mask = attn_mask.bool() | causal_mask.bool()
return attn_mask |