|
|
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 |
|
|
|
|
|
|
|
|
def internvl_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() |
|
|
|
|
|
input_ids, labels = [], [] |
|
|
has_image = any(inst.get('pixel_values') 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' |
|
|
|
|
|
if has_image: |
|
|
pixel_values = [] |
|
|
frames_per_batch = [] |
|
|
|
|
|
for example in instances: |
|
|
input_ids.append(torch.LongTensor(example['input_ids'])) |
|
|
labels.append(torch.LongTensor(example['labels'])) |
|
|
if use_varlen_attn: |
|
|
cumulative_len.append(torch.IntTensor(example['cumulative_len'])) |
|
|
position_ids.append(torch.LongTensor(example['position_ids'])) |
|
|
|
|
|
if has_image: |
|
|
pixel_values.append(example['pixel_values']) |
|
|
|
|
|
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) |
|
|
else: |
|
|
input_ids = torch.stack(input_ids) |
|
|
labels = torch.stack(labels) |
|
|
|
|
|
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: |
|
|
|
|
|
|
|
|
attention_mask = torch.zeros_like(input_ids).bool() |
|
|
for i, length in enumerate(ori_length): |
|
|
attention_mask[i, :length] = True |
|
|
|
|
|
bs, seq_len = input_ids.shape |
|
|
position_ids = torch.arange(seq_len).unsqueeze(0).long().repeat(bs, 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 use_varlen_attn: |
|
|
max_seqlen = ( |
|
|
cumulative_len[0][1:] - |
|
|
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 |
|
|
} |
|
|
else: |
|
|
data_dict = { |
|
|
'input_ids': input_ids, |
|
|
'attention_mask': attention_mask, |
|
|
'position_ids': position_ids, |
|
|
'labels': labels |
|
|
} |
|
|
|
|
|
if has_image: |
|
|
if all(x.shape == pixel_values[0].shape for x in pixel_values): |
|
|
pixel_values = torch.stack(pixel_values, dim=0) |
|
|
data_dict['frames_per_batch'] = frames_per_batch |
|
|
data_dict['pixel_values'] = pixel_values |
|
|
|
|
|
if return_hf_format: |
|
|
return data_dict |
|
|
else: |
|
|
return {'data': data_dict, 'data_samples': None} |
