621 lines
28 KiB
Python
621 lines
28 KiB
Python
# Copyright 2025 Bytedance Ltd. and/or its affiliates.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
|
|
import random
|
|
import json
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from .data_utils import (
|
|
get_flattened_position_ids_interpolate,
|
|
get_flattened_position_ids_extrapolate,
|
|
len2weight,
|
|
patchify,
|
|
prepare_attention_mask_per_sample,
|
|
)
|
|
from .dataset_info import DATASET_INFO, DATASET_REGISTRY
|
|
from .transforms import ImageTransform
|
|
from .video_utils import FrameSampler
|
|
|
|
|
|
class DataConfig:
|
|
def __init__(
|
|
self,
|
|
grouped_datasets,
|
|
text_cond_dropout_prob=0.1,
|
|
vit_cond_dropout_prob=0.4,
|
|
vae_cond_dropout_prob=0.1,
|
|
vae_image_downsample=16,
|
|
max_latent_size=32,
|
|
vit_patch_size=14,
|
|
max_num_patch_per_side=70,
|
|
):
|
|
self.grouped_datasets = grouped_datasets
|
|
self.text_cond_dropout_prob = text_cond_dropout_prob
|
|
self.vit_cond_dropout_prob = vit_cond_dropout_prob
|
|
self.vit_patch_size = vit_patch_size
|
|
self.max_num_patch_per_side = max_num_patch_per_side
|
|
self.vae_cond_dropout_prob = vae_cond_dropout_prob
|
|
self.vae_image_downsample = vae_image_downsample
|
|
self.max_latent_size = max_latent_size
|
|
|
|
|
|
class PackedDataset(torch.utils.data.IterableDataset):
|
|
def __init__(
|
|
self,
|
|
data_config,
|
|
tokenizer,
|
|
special_tokens,
|
|
local_rank,
|
|
world_size,
|
|
num_workers,
|
|
expected_num_tokens=32768,
|
|
max_num_tokens_per_sample=16384,
|
|
max_num_tokens=36864,
|
|
prefer_buffer_before=16384,
|
|
max_buffer_size=50,
|
|
interpolate_pos=False,
|
|
use_flex=False,
|
|
data_status=None,
|
|
):
|
|
super().__init__()
|
|
self.expected_num_tokens = expected_num_tokens
|
|
self.max_num_tokens_per_sample = max_num_tokens_per_sample
|
|
self.prefer_buffer_before = prefer_buffer_before
|
|
self.max_num_tokens = max_num_tokens
|
|
self.max_buffer_size = max_buffer_size
|
|
self.tokenizer = tokenizer
|
|
self.local_rank = local_rank
|
|
self.world_size = world_size
|
|
self.num_workers = num_workers
|
|
self.use_flex = use_flex
|
|
for k, v in special_tokens.items():
|
|
setattr(self, k, v)
|
|
|
|
grouped_datasets, is_mandatory, grouped_weights = self.build_datasets(
|
|
data_config.grouped_datasets, data_status
|
|
)
|
|
self.grouped_datasets = grouped_datasets
|
|
self.dataset_iters = [iter(dataset) for dataset in grouped_datasets]
|
|
self.is_mandatory = is_mandatory
|
|
self.grouped_weights = grouped_weights
|
|
self.data_config = data_config
|
|
self.interpolate_pos = interpolate_pos
|
|
if self.interpolate_pos:
|
|
self.get_flattened_position_ids = get_flattened_position_ids_interpolate
|
|
else:
|
|
self.get_flattened_position_ids = get_flattened_position_ids_extrapolate
|
|
|
|
def build_datasets(self, datasets_metainfo, data_status):
|
|
datasets = []
|
|
is_mandatory = []
|
|
grouped_weights = []
|
|
for grouped_dataset_name, dataset_args in datasets_metainfo.items():
|
|
is_mandatory.append(dataset_args.pop('is_mandatory', False))
|
|
grouped_weights.append(dataset_args.pop('weight', 0.0))
|
|
|
|
if 'frame_sampler_args' in dataset_args.keys():
|
|
frame_sampler = FrameSampler(**dataset_args.pop('frame_sampler_args'))
|
|
dataset_args['frame_sampler'] = frame_sampler
|
|
if 'image_transform_args' in dataset_args.keys():
|
|
transform = ImageTransform(**dataset_args.pop('image_transform_args'))
|
|
dataset_args['transform'] = transform
|
|
if 'vit_image_transform_args' in dataset_args.keys():
|
|
vit_transform = ImageTransform(**dataset_args.pop('vit_image_transform_args'))
|
|
dataset_args['vit_transform'] = vit_transform
|
|
|
|
assert 'dataset_names' in dataset_args.keys()
|
|
dataset_names = dataset_args.pop('dataset_names')
|
|
dataset_args['data_dir_list'] = []
|
|
for item in dataset_names:
|
|
if self.local_rank == 0:
|
|
print(f'Preparing Dataset {grouped_dataset_name}/{item}')
|
|
meta_info = DATASET_INFO[grouped_dataset_name][item]
|
|
dataset_args['data_dir_list'].append(meta_info['data_dir'])
|
|
|
|
if "parquet_info_path" in meta_info.keys():
|
|
if 'parquet_info' not in dataset_args.keys():
|
|
dataset_args['parquet_info'] = {}
|
|
with open(meta_info['parquet_info_path'], 'r') as f:
|
|
parquet_info = json.load(f)
|
|
dataset_args['parquet_info'].update(parquet_info)
|
|
|
|
if 'json_dir' in meta_info.keys():
|
|
# parquet/tar with json
|
|
if 'json_dir_list' not in dataset_args.keys():
|
|
dataset_args['json_dir_list'] = [meta_info['json_dir']]
|
|
else:
|
|
dataset_args['json_dir_list'].append(meta_info['json_dir'])
|
|
|
|
if 'jsonl_path' in meta_info.keys():
|
|
# jsonl with jpeg
|
|
if 'jsonl_path_list' not in dataset_args.keys():
|
|
dataset_args['jsonl_path_list'] = [meta_info['jsonl_path']]
|
|
else:
|
|
dataset_args['jsonl_path_list'].append(meta_info['jsonl_path'])
|
|
|
|
resume_data_status = dataset_args.pop('resume_data_status', True)
|
|
if data_status is not None and grouped_dataset_name in data_status.keys() and resume_data_status:
|
|
data_status_per_group = data_status[grouped_dataset_name]
|
|
else:
|
|
data_status_per_group = None
|
|
dataset = DATASET_REGISTRY[grouped_dataset_name](
|
|
dataset_name=grouped_dataset_name,
|
|
tokenizer=self.tokenizer,
|
|
local_rank=self.local_rank,
|
|
world_size=self.world_size,
|
|
num_workers=self.num_workers,
|
|
data_status=data_status_per_group,
|
|
**dataset_args
|
|
)
|
|
datasets.append(dataset)
|
|
|
|
return datasets, is_mandatory, grouped_weights
|
|
|
|
def set_epoch(self, seed):
|
|
for dataset in self.grouped_datasets:
|
|
dataset.set_epoch(seed)
|
|
|
|
def set_sequence_status(self):
|
|
sequence_status = dict(
|
|
curr = 0,
|
|
sample_lens = list(),
|
|
packed_position_ids = list(),
|
|
nested_attention_masks = list(),
|
|
split_lens = list(),
|
|
attn_modes = list(),
|
|
packed_text_ids = list(),
|
|
packed_text_indexes = list(),
|
|
packed_label_ids = list(),
|
|
ce_loss_indexes = list(),
|
|
ce_loss_weights = list(),
|
|
vae_image_tensors = list(),
|
|
packed_latent_position_ids = list(),
|
|
vae_latent_shapes = list(),
|
|
packed_vae_token_indexes = list(),
|
|
packed_timesteps = list(),
|
|
mse_loss_indexes = list(),
|
|
packed_vit_tokens = list(),
|
|
vit_token_seqlens = list(),
|
|
packed_vit_position_ids = list(),
|
|
packed_vit_token_indexes = list(),
|
|
)
|
|
return sequence_status
|
|
|
|
def to_tensor(self, sequence_status):
|
|
data = dict(
|
|
sequence_length=sum(sequence_status['sample_lens']),
|
|
sample_lens=sequence_status['sample_lens'],
|
|
packed_text_ids=torch.tensor(sequence_status['packed_text_ids']),
|
|
packed_text_indexes=torch.tensor(sequence_status['packed_text_indexes']),
|
|
packed_position_ids=torch.tensor(sequence_status['packed_position_ids']),
|
|
)
|
|
if not self.use_flex:
|
|
data['nested_attention_masks'] = sequence_status['nested_attention_masks']
|
|
else:
|
|
sequence_len = data['sequence_length']
|
|
pad_len = self.max_num_tokens - sequence_len
|
|
data['split_lens'] = sequence_status['split_lens'] + [pad_len]
|
|
data['attn_modes'] = sequence_status['attn_modes'] + ['causal']
|
|
data['sample_lens'] += [pad_len]
|
|
|
|
# if the model has a convnet vae (e.g., as visual tokenizer)
|
|
if len(sequence_status['vae_image_tensors']) > 0:
|
|
image_tensors = sequence_status.pop('vae_image_tensors')
|
|
image_sizes = [item.shape for item in image_tensors]
|
|
max_image_size = [max(item) for item in list(zip(*image_sizes))]
|
|
padded_images = torch.zeros(size=(len(image_tensors), *max_image_size))
|
|
for i, image_tensor in enumerate(image_tensors):
|
|
padded_images[i, :, :image_tensor.shape[1], :image_tensor.shape[2]] = image_tensor
|
|
|
|
data['padded_images'] = padded_images
|
|
data['patchified_vae_latent_shapes'] = sequence_status['vae_latent_shapes']
|
|
data['packed_latent_position_ids'] = torch.cat(sequence_status['packed_latent_position_ids'], dim=0)
|
|
data['packed_vae_token_indexes'] = torch.tensor(sequence_status['packed_vae_token_indexes'])
|
|
|
|
# if the model has a vit (e.g., as visual tokenizer)
|
|
if len(sequence_status['packed_vit_tokens']) > 0:
|
|
data['packed_vit_tokens'] = torch.cat(sequence_status['packed_vit_tokens'], dim=0)
|
|
data['packed_vit_position_ids'] = torch.cat(sequence_status['packed_vit_position_ids'], dim=0)
|
|
data['packed_vit_token_indexes'] = torch.tensor(sequence_status['packed_vit_token_indexes'])
|
|
data['vit_token_seqlens'] = torch.tensor(sequence_status['vit_token_seqlens'])
|
|
|
|
# if the model is required to perform visual generation
|
|
if len(sequence_status['packed_timesteps']) > 0:
|
|
data['packed_timesteps'] = torch.tensor(sequence_status['packed_timesteps'])
|
|
data['mse_loss_indexes'] = torch.tensor(sequence_status['mse_loss_indexes'])
|
|
|
|
# if the model is required to perform text generation
|
|
if len(sequence_status['packed_label_ids']) > 0:
|
|
data['packed_label_ids'] = torch.tensor(sequence_status['packed_label_ids'])
|
|
data['ce_loss_indexes'] = torch.tensor(sequence_status['ce_loss_indexes'])
|
|
data['ce_loss_weights'] = torch.tensor(sequence_status['ce_loss_weights'])
|
|
|
|
return data
|
|
|
|
def __iter__(self):
|
|
total_weights = sum(self.grouped_weights)
|
|
assert total_weights > 0.0
|
|
group_cumprobs = [sum(self.grouped_weights[:i + 1]) / total_weights
|
|
for i in range(len(self.grouped_weights))]
|
|
sequence_status = self.set_sequence_status()
|
|
batch_data_indexes = []
|
|
|
|
buffer = []
|
|
while True:
|
|
# Ensure at least one sample from each group
|
|
if sequence_status['curr'] == 0:
|
|
for group_index, group_iter in enumerate(self.dataset_iters):
|
|
if self.is_mandatory[group_index]:
|
|
while True:
|
|
sample = next(group_iter)
|
|
# if a sample is too long, skip it
|
|
num_tokens = sample['num_tokens'] + 2 * len(sample['sequence_plan'])
|
|
if num_tokens < self.max_num_tokens_per_sample:
|
|
sequence_status = self.pack_sequence(sample, sequence_status)
|
|
batch_data_indexes.append(sample['data_indexes'])
|
|
break
|
|
else:
|
|
print(f"skip a sample with length {num_tokens}")
|
|
continue
|
|
|
|
if sequence_status['curr'] < self.prefer_buffer_before and len(buffer) > 0:
|
|
sample = buffer.pop(0)
|
|
sample_from_buffer = True
|
|
else:
|
|
# sample normally across all groups
|
|
n = random.random()
|
|
group_index = 0
|
|
for i, cumprob in enumerate(group_cumprobs):
|
|
if n < cumprob:
|
|
group_index = i
|
|
break
|
|
sample = next(self.dataset_iters[group_index])
|
|
sample_from_buffer = False
|
|
|
|
# if a sample is too long, skip it
|
|
num_tokens = sample['num_tokens'] + 2 * len(sample['sequence_plan'])
|
|
if num_tokens > self.max_num_tokens_per_sample:
|
|
print(f"skip a sample with length {num_tokens}")
|
|
continue
|
|
|
|
if sequence_status['curr'] + num_tokens > self.max_num_tokens:
|
|
if len(buffer) < self.max_buffer_size and not sample_from_buffer:
|
|
buffer.append(sample)
|
|
else:
|
|
print(f"Yielding data with length {sum(sequence_status['sample_lens'])}")
|
|
data = self.to_tensor(sequence_status)
|
|
data['batch_data_indexes'] = batch_data_indexes
|
|
yield data
|
|
sequence_status = self.set_sequence_status()
|
|
batch_data_indexes = []
|
|
continue
|
|
|
|
sequence_status = self.pack_sequence(sample, sequence_status)
|
|
batch_data_indexes.append(sample['data_indexes'])
|
|
|
|
if sequence_status['curr'] >= self.expected_num_tokens:
|
|
data = self.to_tensor(sequence_status)
|
|
data['batch_data_indexes'] = batch_data_indexes
|
|
yield data
|
|
sequence_status = self.set_sequence_status()
|
|
batch_data_indexes = []
|
|
|
|
def pack_sequence(self, sample, sequence_status):
|
|
image_tensor_list = sample['image_tensor_list']
|
|
text_ids_list = sample['text_ids_list']
|
|
sequence_plan = sample['sequence_plan']
|
|
|
|
split_lens, attn_modes = list(), list()
|
|
curr = sequence_status['curr']
|
|
curr_rope_id = 0
|
|
sample_lens = 0
|
|
|
|
for item in sequence_plan:
|
|
split_start = item.get('split_start', True)
|
|
if split_start:
|
|
curr_split_len = 0
|
|
|
|
if item['type'] == 'text':
|
|
text_ids = text_ids_list.pop(0)
|
|
if item['enable_cfg'] == 1 and random.random() < self.data_config.text_cond_dropout_prob:
|
|
continue
|
|
|
|
shifted_text_ids = [self.bos_token_id] + text_ids
|
|
sequence_status['packed_text_ids'].extend(shifted_text_ids)
|
|
sequence_status['packed_text_indexes'].extend(range(curr, curr + len(shifted_text_ids)))
|
|
if item['loss'] == 1:
|
|
sequence_status['ce_loss_indexes'].extend(range(curr, curr + len(shifted_text_ids)))
|
|
sequence_status['ce_loss_weights'].extend(
|
|
[len2weight(len(shifted_text_ids))] * len(shifted_text_ids)
|
|
)
|
|
sequence_status['packed_label_ids'].extend(text_ids + [self.eos_token_id])
|
|
curr += len(shifted_text_ids)
|
|
curr_split_len += len(shifted_text_ids)
|
|
|
|
# add a <|im_end|> token
|
|
sequence_status['packed_text_ids'].append(self.eos_token_id)
|
|
sequence_status['packed_text_indexes'].append(curr)
|
|
if item['special_token_loss'] == 1: # <|im_end|> may have loss
|
|
sequence_status['ce_loss_indexes'].append(curr)
|
|
sequence_status['ce_loss_weights'].append(1.0)
|
|
sequence_status['packed_label_ids'].append(item['special_token_label'])
|
|
curr += 1
|
|
curr_split_len += 1
|
|
|
|
# update sequence status
|
|
attn_modes.append("causal")
|
|
sequence_status['packed_position_ids'].extend(range(curr_rope_id, curr_rope_id + curr_split_len))
|
|
curr_rope_id += curr_split_len
|
|
|
|
elif item['type'] == 'vit_image':
|
|
image_tensor = image_tensor_list.pop(0)
|
|
if item['enable_cfg'] == 1 and random.random() < self.data_config.vit_cond_dropout_prob:
|
|
curr_rope_id += 1
|
|
continue
|
|
|
|
# add a <|startofimage|> token
|
|
sequence_status['packed_text_ids'].append(self.start_of_image)
|
|
sequence_status['packed_text_indexes'].append(curr)
|
|
curr += 1
|
|
curr_split_len += 1
|
|
|
|
# preprocess image
|
|
vit_tokens = patchify(image_tensor, self.data_config.vit_patch_size)
|
|
num_img_tokens = vit_tokens.shape[0]
|
|
sequence_status['packed_vit_token_indexes'].extend(range(curr, curr + num_img_tokens))
|
|
curr += num_img_tokens
|
|
curr_split_len += num_img_tokens
|
|
|
|
sequence_status['packed_vit_tokens'].append(vit_tokens)
|
|
sequence_status['vit_token_seqlens'].append(num_img_tokens)
|
|
sequence_status['packed_vit_position_ids'].append(
|
|
self.get_flattened_position_ids(
|
|
image_tensor.size(1), image_tensor.size(2),
|
|
self.data_config.vit_patch_size,
|
|
max_num_patches_per_side=self.data_config.max_num_patch_per_side
|
|
)
|
|
)
|
|
|
|
# add a <|endofimage|> token
|
|
sequence_status['packed_text_ids'].append(self.end_of_image)
|
|
sequence_status['packed_text_indexes'].append(curr)
|
|
if item['special_token_loss'] == 1: # <|endofimage|> may have loss
|
|
sequence_status['ce_loss_indexes'].append(curr)
|
|
sequence_status['ce_loss_weights'].append(1.0)
|
|
sequence_status['packed_label_ids'].append(item['special_token_label'])
|
|
curr += 1
|
|
curr_split_len += 1
|
|
|
|
# update sequence status
|
|
attn_modes.append("full")
|
|
sequence_status['packed_position_ids'].extend([curr_rope_id] * curr_split_len)
|
|
curr_rope_id += 1
|
|
|
|
elif item['type'] == 'vae_image':
|
|
image_tensor = image_tensor_list.pop(0)
|
|
if item['enable_cfg'] == 1 and random.random() < self.data_config.vae_cond_dropout_prob:
|
|
# FIXME fix vae dropout in video2video setting.
|
|
curr_rope_id += 1
|
|
continue
|
|
|
|
# add a <|startofimage|> token
|
|
sequence_status['packed_text_ids'].append(self.start_of_image)
|
|
sequence_status['packed_text_indexes'].append(curr)
|
|
curr += 1
|
|
curr_split_len += 1
|
|
|
|
# preprocess image
|
|
sequence_status['vae_image_tensors'].append(image_tensor)
|
|
sequence_status['packed_latent_position_ids'].append(
|
|
self.get_flattened_position_ids(
|
|
image_tensor.size(1), image_tensor.size(2),
|
|
self.data_config.vae_image_downsample,
|
|
max_num_patches_per_side=self.data_config.max_latent_size
|
|
)
|
|
)
|
|
H, W = image_tensor.shape[1:]
|
|
h = H // self.data_config.vae_image_downsample
|
|
w = W // self.data_config.vae_image_downsample
|
|
sequence_status['vae_latent_shapes'].append((h, w))
|
|
|
|
num_img_tokens = w * h
|
|
sequence_status['packed_vae_token_indexes'].extend(range(curr, curr + num_img_tokens))
|
|
if item['loss'] == 1:
|
|
sequence_status['mse_loss_indexes'].extend(range(curr, curr + num_img_tokens))
|
|
if split_start:
|
|
timestep = np.random.randn()
|
|
else:
|
|
timestep = float('-inf')
|
|
|
|
sequence_status['packed_timesteps'].extend([timestep] * num_img_tokens)
|
|
curr += num_img_tokens
|
|
curr_split_len += num_img_tokens
|
|
|
|
# add a <|endofimage|> token
|
|
sequence_status['packed_text_ids'].append(self.end_of_image)
|
|
sequence_status['packed_text_indexes'].append(curr)
|
|
# <|endofimage|> may have loss
|
|
if item['special_token_loss'] == 1:
|
|
sequence_status['ce_loss_indexes'].append(curr)
|
|
sequence_status['ce_loss_weights'].append(1.0)
|
|
sequence_status['packed_label_ids'].append(item['special_token_label'])
|
|
curr += 1
|
|
curr_split_len += 1
|
|
|
|
# update sequence status
|
|
if split_start:
|
|
if item['loss'] == 1 and 'frame_delta' not in item.keys():
|
|
attn_modes.append("noise")
|
|
else:
|
|
attn_modes.append("full")
|
|
sequence_status['packed_position_ids'].extend([curr_rope_id] * (num_img_tokens + 2))
|
|
if 'frame_delta' in item.keys():
|
|
curr_rope_id += item['frame_delta']
|
|
elif item['loss'] == 0:
|
|
curr_rope_id += 1
|
|
|
|
if item.get('split_end', True):
|
|
split_lens.append(curr_split_len)
|
|
sample_lens += curr_split_len
|
|
|
|
sequence_status['curr'] = curr
|
|
sequence_status['sample_lens'].append(sample_lens)
|
|
# prepare attention mask
|
|
if not self.use_flex:
|
|
sequence_status['nested_attention_masks'].append(
|
|
prepare_attention_mask_per_sample(split_lens, attn_modes)
|
|
)
|
|
else:
|
|
sequence_status['split_lens'].extend(split_lens)
|
|
sequence_status['attn_modes'].extend(attn_modes)
|
|
|
|
return sequence_status
|
|
|
|
|
|
class SimpleCustomBatch:
|
|
def __init__(self, batch):
|
|
data = batch[0]
|
|
self.batch_data_indexes = data['batch_data_indexes']
|
|
self.sequence_length = data["sequence_length"]
|
|
self.sample_lens = data["sample_lens"]
|
|
self.packed_text_ids = data["packed_text_ids"]
|
|
self.packed_text_indexes = data["packed_text_indexes"]
|
|
self.packed_position_ids = data["packed_position_ids"]
|
|
|
|
self.use_flex = "nested_attention_masks" not in data.keys()
|
|
|
|
if self.use_flex:
|
|
self.split_lens = data["split_lens"]
|
|
self.attn_modes = data["attn_modes"]
|
|
else:
|
|
self.nested_attention_masks = data["nested_attention_masks"]
|
|
|
|
if "padded_images" in data.keys():
|
|
self.padded_images = data["padded_images"]
|
|
self.patchified_vae_latent_shapes = data["patchified_vae_latent_shapes"]
|
|
self.packed_latent_position_ids = data["packed_latent_position_ids"]
|
|
self.packed_vae_token_indexes = data["packed_vae_token_indexes"]
|
|
|
|
if "packed_vit_tokens" in data.keys():
|
|
self.packed_vit_tokens = data["packed_vit_tokens"]
|
|
self.packed_vit_position_ids = data["packed_vit_position_ids"]
|
|
self.packed_vit_token_indexes = data["packed_vit_token_indexes"]
|
|
self.vit_token_seqlens = data["vit_token_seqlens"]
|
|
|
|
if "packed_timesteps" in data.keys():
|
|
self.packed_timesteps = data["packed_timesteps"]
|
|
self.mse_loss_indexes = data["mse_loss_indexes"]
|
|
|
|
if "packed_label_ids" in data.keys():
|
|
self.packed_label_ids = data["packed_label_ids"]
|
|
self.ce_loss_indexes = data["ce_loss_indexes"]
|
|
self.ce_loss_weights = data["ce_loss_weights"]
|
|
|
|
def pin_memory(self):
|
|
self.packed_text_ids = self.packed_text_ids.pin_memory()
|
|
self.packed_text_indexes = self.packed_text_indexes.pin_memory()
|
|
self.packed_position_ids = self.packed_position_ids.pin_memory()
|
|
|
|
if not self.use_flex:
|
|
self.nested_attention_masks = [item.pin_memory() for item in self.nested_attention_masks]
|
|
|
|
if hasattr(self, 'padded_images'):
|
|
self.padded_images = self.padded_images.pin_memory()
|
|
self.packed_vae_token_indexes = self.packed_vae_token_indexes.pin_memory()
|
|
self.packed_latent_position_ids = self.packed_latent_position_ids.pin_memory()
|
|
|
|
if hasattr(self, 'packed_timesteps'):
|
|
self.packed_timesteps = self.packed_timesteps.pin_memory()
|
|
self.mse_loss_indexes = self.mse_loss_indexes.pin_memory()
|
|
|
|
if hasattr(self, 'packed_vit_tokens'):
|
|
self.packed_vit_tokens = self.packed_vit_tokens.pin_memory()
|
|
self.packed_vit_position_ids = self.packed_vit_position_ids.pin_memory()
|
|
self.packed_vit_token_indexes = self.packed_vit_token_indexes.pin_memory()
|
|
self.vit_token_seqlens = self.vit_token_seqlens.pin_memory()
|
|
|
|
if hasattr(self, 'packed_label_ids'):
|
|
self.packed_label_ids = self.packed_label_ids.pin_memory()
|
|
self.ce_loss_indexes = self.ce_loss_indexes.pin_memory()
|
|
self.ce_loss_weights = self.ce_loss_weights.pin_memory()
|
|
|
|
return self
|
|
|
|
def cuda(self, device):
|
|
self.packed_text_ids = self.packed_text_ids.to(device)
|
|
self.packed_text_indexes = self.packed_text_indexes.to(device)
|
|
self.packed_position_ids = self.packed_position_ids.to(device)
|
|
|
|
if not self.use_flex:
|
|
self.nested_attention_masks = [item.to(device) for item in self.nested_attention_masks]
|
|
|
|
if hasattr(self, 'padded_images'):
|
|
self.padded_images = self.padded_images.to(device)
|
|
self.packed_vae_token_indexes = self.packed_vae_token_indexes.to(device)
|
|
self.packed_latent_position_ids = self.packed_latent_position_ids.to(device)
|
|
|
|
if hasattr(self, 'packed_timesteps'):
|
|
self.packed_timesteps = self.packed_timesteps.to(device)
|
|
self.mse_loss_indexes = self.mse_loss_indexes.to(device)
|
|
|
|
if hasattr(self, 'packed_vit_tokens'):
|
|
self.packed_vit_tokens = self.packed_vit_tokens.to(device)
|
|
self.packed_vit_position_ids = self.packed_vit_position_ids.to(device)
|
|
self.packed_vit_token_indexes = self.packed_vit_token_indexes.to(device)
|
|
self.vit_token_seqlens = self.vit_token_seqlens.to(device)
|
|
|
|
if hasattr(self, 'packed_label_ids'):
|
|
self.packed_label_ids = self.packed_label_ids.to(device)
|
|
self.ce_loss_indexes = self.ce_loss_indexes.to(device)
|
|
self.ce_loss_weights = self.ce_loss_weights.to(device)
|
|
|
|
return self
|
|
|
|
def to_dict(self):
|
|
data = dict(
|
|
sequence_length = self.sequence_length,
|
|
sample_lens = self.sample_lens,
|
|
packed_text_ids = self.packed_text_ids,
|
|
packed_text_indexes = self.packed_text_indexes,
|
|
packed_position_ids = self.packed_position_ids,
|
|
batch_data_indexes = self.batch_data_indexes,
|
|
)
|
|
|
|
if not self.use_flex:
|
|
data['nested_attention_masks'] = self.nested_attention_masks
|
|
else:
|
|
data['split_lens'] = self.split_lens
|
|
data['attn_modes'] = self.attn_modes
|
|
|
|
if hasattr(self, 'padded_images'):
|
|
data['padded_images'] = self.padded_images
|
|
data['patchified_vae_latent_shapes'] = self.patchified_vae_latent_shapes
|
|
data['packed_latent_position_ids'] = self.packed_latent_position_ids
|
|
data['packed_vae_token_indexes'] = self.packed_vae_token_indexes
|
|
|
|
if hasattr(self, 'packed_vit_tokens'):
|
|
data['packed_vit_tokens'] = self.packed_vit_tokens
|
|
data['packed_vit_position_ids'] = self.packed_vit_position_ids
|
|
data['packed_vit_token_indexes'] = self.packed_vit_token_indexes
|
|
data['vit_token_seqlens'] = self.vit_token_seqlens
|
|
|
|
if hasattr(self, 'packed_timesteps'):
|
|
data['packed_timesteps'] = self.packed_timesteps
|
|
data['mse_loss_indexes'] = self.mse_loss_indexes
|
|
|
|
if hasattr(self, 'packed_label_ids'):
|
|
data['packed_label_ids'] = self.packed_label_ids
|
|
data['ce_loss_indexes'] = self.ce_loss_indexes
|
|
data['ce_loss_weights'] = self.ce_loss_weights
|
|
|
|
return data
|
|
|
|
|
|
def collate_wrapper():
|
|
def collate_fn(batch):
|
|
return SimpleCustomBatch(batch)
|
|
return collate_fn
|