# Copyright 2025 MMaDA Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
import sys
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
os.environ["TOKENIZERS_PARALLELISM"] = "true"
import json
import logging
import math
import shutil
import time
from pathlib import Path
from typing import Union
import numpy as np
from PIL import Image
from omegaconf import OmegaConf
import wandb
import torch
from torch.optim import AdamW
from lightning.pytorch.utilities import CombinedLoader
from transformers import AutoTokenizer, AutoConfig
from accelerate import Accelerator
from accelerate.logging import get_logger
from accelerate.utils import DistributedType, set_seed
from training.data import Text2ImageDataset
from training.utils import get_config, flatten_omega_conf, image_transform
from training.imagenet_dataset import ImageNetDataset
from parquet import RefinedWebDataset
from models import MAGVITv2, get_mask_schedule, MMadaModelLM, MMadaConfig
from training.prompting_utils import UniversalPrompting
from models.lr_schedulers import get_scheduler
from models.logging import set_verbosity_info, set_verbosity_error
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler
SYSTEM_PROMPT_LEN = 28
from training.utils import get_config, flatten_omega_conf, mask_or_random_replace_tokens, AverageMeter
try:
import apex
is_apex_available = True
except ImportError:
is_apex_available = False
logger = get_logger(__name__, log_level="INFO")
def get_vq_model_class(model_type):
if model_type == "magvitv2":
return MAGVITv2
elif model_type == "vq16":
return VQ_16
else:
raise ValueError(f"model_type {model_type} not supported.")
def main():
#########################
# SETUP Accelerator #
#########################
config = get_config()
# Enable TF32 on Ampere GPUs
if config.training.enable_tf32:
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.benchmark = True
torch.backends.cudnn.deterministic = False
config.experiment.logging_dir = str(Path(config.experiment.output_dir) / "logs")
accelerator = Accelerator(
gradient_accumulation_steps=config.training.gradient_accumulation_steps,
mixed_precision=config.training.mixed_precision,
log_with="wandb",
project_dir=config.experiment.logging_dir,
split_batches=True,
)
total_batch_size_per_gpu = (config.training.batch_size_t2i
+ config.training.batch_size_lm
+ config.training.batch_size_mmu)
total_batch_size = (
(config.training.batch_size_t2i + config.training.batch_size_lm + config.training.batch_size_mmu)
* accelerator.num_processes * config.training.gradient_accumulation_steps
)
if accelerator.distributed_type == DistributedType.DEEPSPEED:
accelerator.state.deepspeed_plugin.deepspeed_config["train_micro_batch_size_per_gpu"] = (
total_batch_size_per_gpu
)
#####################################
# SETUP LOGGING, SEED and CONFIG #
#####################################
# Make one log on every process with the configuration for debugging.
logging.basicConfig(
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
datefmt="%m/%d/%Y %H:%M:%S",
level=logging.INFO,
)
logger.info(accelerator.state, main_process_only=False)
if accelerator.is_local_main_process:
set_verbosity_info()
else:
set_verbosity_error()
# We need to initialize the trackers we use, and also store our configuration.
# The trackers initializes automatically on the main process.
if accelerator.is_main_process:
resume_wandb_run = config.wandb.resume
run_id = config.wandb.get("run_id", None)
if run_id is None:
resume_wandb_run = False
run_id = wandb.util.generate_id()
config.wandb.run_id = run_id
wandb_init_kwargs = dict(
name=config.experiment.name,
id=run_id,
resume=resume_wandb_run,
entity=config.wandb.get("entity", None),
config_exclude_keys=[],
)
wandb_config = {k: v for k, v in flatten_omega_conf(config, resolve=True)}
wandb_config.pop("experiment.resume_from_checkpoint")
accelerator.init_trackers(
config.experiment.project,
config=wandb_config,
init_kwargs={"wandb": wandb_init_kwargs},
)
if accelerator.is_main_process:
os.makedirs(config.experiment.output_dir, exist_ok=True)
config_path = Path(config.experiment.output_dir) / "config.yaml"
logging.info(f"Saving config to {config_path}")
OmegaConf.save(config, config_path)
# If passed along, set the training seed now.
if config.training.seed is not None:
set_seed(config.training.seed)
#########################
# MODELS and OPTIMIZER #
#########################
logger.info("Loading models and optimizer")
tokenizer = AutoTokenizer.from_pretrained(config.model.mmada.pretrained_model_path, padding_side="left")
uni_prompting = UniversalPrompting(tokenizer, max_text_len=config.dataset.preprocessing.max_seq_length,
special_tokens=(
"<|soi|>", "<|eoi|>", "<|sov|>", "<|eov|>", "<|t2i|>",
"<|mmu|>", "<|t2v|>", "<|v2v|>", "<|lvg|>"
),
ignore_id=-100, cond_dropout_prob=config.training.cond_dropout_prob, use_reserved_token=True)
print('special tokens : \n', uni_prompting.sptids_dict)
# VQ model for processing image into discrete tokens
vq_model = get_vq_model_class(config.model.vq_model.type)
if config.model.vq_model.get("pretrained_model_path", None):
vq_model = vq_model().to(accelerator.device)
state_dict = torch.load(config.model.vq_model.pretrained_model_path)['model']
vq_model.load_state_dict(state_dict)
else:
vq_model = vq_model.from_pretrained(config.model.vq_model.vq_model_name).to(accelerator.device)
vq_model.eval()
vq_model.requires_grad_(False)
# Initialize mmada in pretraining stage
base_config = AutoConfig.from_pretrained(config.model.mmada.pretrained_model_path).to_dict()
mmada_config_dict = {k: v for k, v in config.model.mmada.items()}
merged_config = {**base_config, **mmada_config_dict}
mmada_config = MMadaConfig(**merged_config)
model = MMadaModelLM.from_pretrained(config.model.mmada.pretrained_model_path, torch_dtype=torch.bfloat16, config=mmada_config)
model.resize_token_embeddings(mmada_config.new_vocab_size)
model.config.embedding_size = model.config.vocab_size
model = model.to(accelerator.device)
mask_id = model.config.mask_token_id
##################################
# Optimizer and LR scheduler #
#################################
optimizer_config = config.optimizer.params
# no decay on bias and layernorm and embedding
no_decay = ["bias", "layer_norm.weight", "mlm_ln.weight", "embeddings.weight"]
optimizer_grouped_parameters = [
{
"params": [p for n, p in model.named_parameters() if
p.requires_grad and not any(nd in n for nd in no_decay)],
"weight_decay": optimizer_config.weight_decay,
},
{
"params": [p for n, p in model.named_parameters() if
p.requires_grad and any(nd in n for nd in no_decay)],
"weight_decay": 0.0,
},
]
optimizer_type = config.optimizer.name
if optimizer_type == "adamw":
optimizer = AdamW(
optimizer_grouped_parameters,
lr=optimizer_config.learning_rate,
betas=(optimizer_config.beta1, optimizer_config.beta2),
weight_decay=optimizer_config.weight_decay,
eps=optimizer_config.epsilon,
)
else:
raise ValueError(f"Optimizer {optimizer_type} not supported")
# Create mask scheduler
if config.get("mask_schedule", None) is not None:
schedule = config.mask_schedule.schedule
args = config.mask_schedule.get("params", {})
mask_schedule = get_mask_schedule(schedule, **args)
else:
mask_schedule = get_mask_schedule(config.training.get("mask_schedule", "cosine"))
lr_scheduler = get_scheduler(
config.lr_scheduler.scheduler,
optimizer=optimizer,
num_training_steps=config.training.max_train_steps,
num_warmup_steps=config.lr_scheduler.params.warmup_steps,
min_lr_scale=config.lr_scheduler.params.min_lr_scale
)
##################################
# DATALOADER #
#################################
logger.info("Creating dataloaders and lr_scheduler")
total_batch_size_t2i_without_accum = config.training.batch_size_t2i * accelerator.num_processes
total_batch_size_t2i = (
config.training.batch_size_t2i * accelerator.num_processes * config.training.gradient_accumulation_steps
)
# DataLoaders creation:
# We use webdataset for data loading. The dataloaders are created with sampling with replacement.
# We don't do dataset resuming here, instead we resample the shards and buffer each time. The sampling is stochastic.
# This means that the dataloading is not deterministic, but it's fast and efficient.
preproc_config = config.dataset.preprocessing
dataset_config = config.dataset.params
# Data for generation
if config.dataset.gen_type == "t2i":
dataset = Text2ImageDataset(
train_shards_path_or_url=dataset_config.train_t2i_shards_path_or_url,
tokenizer=None, # we want to get raw texts
max_seq_length=preproc_config.max_seq_length,
num_train_examples=config.experiment.max_train_examples_t2i,
per_gpu_batch_size=config.training.batch_size_t2i,
global_batch_size=total_batch_size_t2i_without_accum,
num_workers=dataset_config.num_workers,
resolution=preproc_config.resolution,
shuffle_buffer_size=dataset_config.shuffle_buffer_size,
pin_memory=dataset_config.pin_memory,
persistent_workers=dataset_config.persistent_workers,
external_caption_path=dataset_config.external_caption_path,
external_journeydb_caption_path=dataset_config.external_journeydb_caption_path,
external_laion12m_caption_path=dataset_config.external_laion12m_caption_path,
external_cc12m_caption_path=dataset_config.external_cc12m_caption_path,
)
train_dataloader_t2i = dataset.train_dataloader
num_update_steps_per_epoch = math.ceil(
train_dataloader_t2i.num_batches / config.training.gradient_accumulation_steps)
num_train_epochs = math.ceil(config.training.max_train_steps / num_update_steps_per_epoch)
elif config.dataset.gen_type == "t2i_parquet":
# this part relies on the internal packages, which will not be released
num_update_steps_per_epoch = math.ceil(config.experiment.max_train_examples_t2i / total_batch_size_t2i)
num_train_epochs = math.ceil(config.training.max_train_steps / num_update_steps_per_epoch)
train_dataloader_t2i = create_imagetext_dataloader(
train_shards_path_or_url=dataset_config.train_t2i_shards_path_or_url,
batch_size=config.training.batch_size_t2i,
image_size=preproc_config.resolution,
num_workers=dataset_config.num_workers,
num_readers=32,
predefined_steps=num_update_steps_per_epoch,
drop_last=True,
shuffle=True,
shuffle_buffer_size=dataset_config.shuffle_buffer_size
)
elif config.dataset.gen_type == "imagenet1k":
dataset_imagenet = ImageNetDataset(
dataset_config.train_t2i_shards_path_or_url,
image_size=preproc_config.resolution,
)
print('process index : ',
accelerator.process_index, ', ', accelerator.num_processes,
"Length: ", len(dataset_imagenet))
if accelerator.num_processes > 1:
sampler = DistributedSampler(dataset_imagenet,
num_replicas=accelerator.num_processes,
rank=accelerator.process_index,
shuffle=True,
)
shuffle = False
else:
sampler = None
shuffle = True
train_dataloader_t2i = DataLoader(dataset_imagenet, batch_size=config.training.batch_size_t2i,
sampler=sampler, collate_fn=dataset_imagenet.collate_fn,
shuffle=shuffle, num_workers=dataset_config.num_workers)
num_update_steps_per_epoch = math.ceil(len(dataset_imagenet) / total_batch_size_t2i)
num_train_epochs = math.ceil(config.training.max_train_steps / num_update_steps_per_epoch)
else:
raise ValueError(f"Unsupported dataset type {config.dataset.type}")
total_batch_size_mmu_without_accum = config.training.batch_size_mmu * accelerator.num_processes
# Data for image captioning
if config.dataset.und_type == "captioning":
dataset_mmu = Text2ImageDataset(
train_shards_path_or_url=dataset_config.train_mmu_shards_path_or_url,
tokenizer=None, # we want to get raw texts
max_seq_length=preproc_config.max_seq_length,
num_train_examples=config.experiment.max_train_examples_mmu,
per_gpu_batch_size=config.training.batch_size_mmu,
global_batch_size=total_batch_size_mmu_without_accum,
num_workers=dataset_config.num_workers,
resolution=preproc_config.resolution,
shuffle_buffer_size=dataset_config.shuffle_buffer_size,
pin_memory=dataset_config.pin_memory,
persistent_workers=dataset_config.persistent_workers,
external_caption_path=dataset_config.external_caption_path,
external_journeydb_caption_path=dataset_config.external_journeydb_caption_path,
external_laion12m_caption_path=dataset_config.external_laion12m_caption_path,
external_cc12m_caption_path=dataset_config.external_cc12m_caption_path,
is_captioning=True,
add_caption_prompt=dataset_config.add_caption_prompt,
)
train_dataloader_mmu = dataset_mmu.train_dataloader
elif config.dataset.und_type == "captioning_parquet":
train_dataloader_mmu = create_imagetext_dataloader(
train_shards_path_or_url=dataset_config.train_mmu_shards_path_or_url,
batch_size=config.training.batch_size_mmu,
image_size=preproc_config.resolution,
num_workers=dataset_config.num_workers,
num_readers=32,
predefined_steps=num_update_steps_per_epoch,
drop_last=True,
shuffle=True,
shuffle_buffer_size=dataset_config.shuffle_buffer_size,
is_captioning=True
)
else:
raise NotImplementedError(f"Unsupported dataset type {config.dataset.und_type}")
# LLM pure text dataset: RefinedWeb
dataset_lm = RefinedWebDataset(data_path=dataset_config.train_lm_shards_path_or_url,
rank=accelerator.process_index,
world_size=accelerator.num_processes,
num_workers=dataset_config.num_workers)
train_dataloader_lm = torch.utils.data.DataLoader(dataset_lm, batch_size=config.training.batch_size_lm,
sampler=None, collate_fn=dataset_lm.collate_fn,
num_workers=dataset_config.num_workers)
# Combine these dataloaders into a single iterable model
iterables = {
"t2i_flow": train_dataloader_t2i,
"lm_flow": train_dataloader_lm,
"mmu_flow": train_dataloader_mmu,
}
combined_dataloader = CombinedLoader(iterables, mode=config.dataset.combined_loader_mode)
##################################
# MODEL RESUME #
#################################
global_step = 0
first_epoch = 0
if config.experiment.resume_from_checkpoint:
dirs = os.listdir(config.experiment.output_dir)
logger.info(f"dirs: {dirs}")
dirs = [d for d in dirs if d.startswith("checkpoint")]
dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
path = dirs[-1] if len(dirs) > 0 else None
logger.info(f"path: {path}")
if path is not None:
path = os.path.join(config.experiment.output_dir, path)
logger.info(f"Resuming from checkpoint: {path}")
global_step = int(os.path.basename(path).split("-")[1])
first_epoch = global_step // num_update_steps_per_epoch
if os.path.exists(f'{path}/unwrapped_model/pytorch_model.bin'):
state_dict = torch.load(f'{path}/unwrapped_model/pytorch_model.bin', map_location="cpu")
model.load_state_dict(state_dict, strict=True)
del state_dict
elif os.path.exists(f'{path}/unwrapped_model/pytorch_model.bin.index.json'):
from safetensors.torch import load_file
from transformers.modeling_utils import load_sharded_checkpoint
load_sharded_checkpoint(model, f'{path}/unwrapped_model/')
# if safetensors sharded checkpoint exists
elif os.path.exists(f'{path}/unwrapped_model/model.safetensors.index.json'):
from transformers.modeling_utils import load_sharded_checkpoint
load_sharded_checkpoint(
model,
f'{path}/unwrapped_model/',
# weight_map=None,
# load_state_dict_fn="safetensors"
)
else:
raise FileNotFoundError(f"Checkpoint {path}/unwrapped_model/pytorch_model.bin not found")
else:
logger.info("Not resuming from checkpoint")
##################################
# Prepare accelerator #
#################################
logger.info("Preparing model, optimizer and dataloaders")
model, optimizer, lr_scheduler = accelerator.prepare(model, optimizer, lr_scheduler)
vq_model.to(device=accelerator.device)
mask_dtype = model.get_input_embeddings().weight.dtype
##################################
# Training #
#################################
logger.info("***** Running training *****")
logger.info(f" Num training steps = {config.training.max_train_steps}")
logger.info(f" Instantaneous batch size per device = {total_batch_size_per_gpu}")
logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
logger.info(f" Gradient Accumulation steps = {config.training.gradient_accumulation_steps}")
@torch.no_grad()
def prepare_inputs_and_labels(
pixel_values_or_image_ids: Union[torch.FloatTensor, torch.LongTensor],
texts: Union[str, str],
min_masking_rate: float = 0.0,
is_train: bool = True,
):
image_tokens = vq_model.get_code(pixel_values_or_image_ids)
image_tokens = image_tokens + len(uni_prompting.text_tokenizer)
# create MLM mask and labels
input_ids, labels, loss_weight, mask_prob = mask_or_random_replace_tokens(
image_tokens,
mask_id,
config,
mask_schedule=mask_schedule,
is_train=is_train,
)
input_ids, masks, labels = uni_prompting((texts, input_ids, labels), 't2i')
return input_ids, labels, mask_prob, image_tokens, masks
@torch.no_grad()
def prepare_inputs_and_labels_for_text(
texts: Union[str, str], max_seq_len, eps=1e-3
):
# create MLM mask and labels
input_ids_lm, prompt_mask, labels_lm = uni_prompting((texts_lm, max_seq_len), 'lm')
b, l = input_ids_lm.shape
t = torch.rand(b, device=input_ids_lm.device)
p_mask = (1 - eps) * t + eps
p_mask = p_mask[:, None].repeat(1, l)
masked_indices = torch.rand((b, l), device=input_ids_lm.device) < p_mask
# 126336 is used for [MASK] token
noisy_batch = torch.where(masked_indices, mask_id, input_ids_lm)
masked_indices = noisy_batch == mask_id
return noisy_batch, labels_lm, p_mask
@torch.no_grad()
def prepare_inputs_and_labels_for_mmu(
input_ids_mmu, prompt_masks, labels_mmu, eps=1e-3
):
b, l = input_ids_mmu.shape
t = torch.rand(b, device=input_ids_mmu.device)
p_mask = (1 - eps) * t + eps
p_mask = p_mask[:, None].repeat(1, l)
masked_indices = torch.rand((b, l), device=input_ids_mmu.device) < p_mask
# 126336 is used for [MASK] token
noisy_batch = torch.where(masked_indices, mask_id, input_ids_mmu)
masked_indices = noisy_batch == mask_id
noisy_batch[prompt_masks.bool()] = input_ids_mmu[prompt_masks.bool()]
masked_indices = noisy_batch == mask_id
prompt_masks = prompt_masks.to(torch.int64)
answer_lengths = torch.sum((1 - prompt_masks), dim=-1, keepdim=True)
answer_lengths = answer_lengths.repeat(1, noisy_batch.shape[1])
return noisy_batch, labels_mmu, p_mask, answer_lengths
batch_time_m = AverageMeter()
data_time_m = AverageMeter()
end = time.time()
for epoch in range(first_epoch, num_train_epochs):
model.train()
for batch, batch_idx, dataloader_idx in combined_dataloader:
# for loss calculation
batch_size_t2i = batch["t2i_flow"]["images"].shape[0]
batch_size_lm = len(batch["lm_flow"]["input_ids"])
batch_size_mmu = batch["mmu_flow"]["images"].shape[0]
# *-------*-------*-------*-------*-------*-------*-------*-------*-------*-------*-------*
# Build formatted sequences for class-conditional/text-to-image generation
# *-------*-------*-------*-------*-------*-------*-------*-------*-------*-------*-------*
pixel_values, texts = batch["t2i_flow"]["images"], batch["t2i_flow"]["input_ids"]
pixel_values = pixel_values.to(accelerator.device, non_blocking=True)
data_time_m.update(time.time() - end)
# Encode images to image tokens, mask them and create input and labels
(
input_ids,
labels,
mask_prob,
image_tokens_ori,
t2i_masks
) = prepare_inputs_and_labels(pixel_values, texts, config.training.min_masking_rate)
# *-------*-------*-------*-------*-------*-------*-------*-------*-------*-------*-------*
# Build formatted sequences for language modeling
# *-------*-------*-------*-------*-------*-------*-------*-------*-------*-------*-------*
max_seq_len = input_ids.shape[-1]
texts_lm = batch["lm_flow"]["input_ids"]
(
input_ids_lm,
labels_lm,
p_mask_lm
) = prepare_inputs_and_labels_for_text(texts_lm, max_seq_len)
input_ids = torch.cat((input_ids, input_ids_lm.to(input_ids.device)), dim=0)
labels = torch.cat((labels, labels_lm.to(input_ids.device)), dim=0)
# *-------*-------*-------*-------*-------*-------*-------*-------*-------*-------*-------*
# Build formatted sequences for captioning/multimodal understanding
# *-------*-------*-------*-------*-------*-------*-------*-------*-------*-------*-------*
if "llava" in config.dataset.und_type:
pixel_values_mmu, input_ids_mmu, labels_mmu = (batch["mmu_flow"]["images"], batch["mmu_flow"]["input_ids"],batch["mmu_flow"]["labels"])
pixel_values_mmu = pixel_values_mmu.to(accelerator.device, non_blocking=True)
input_ids_mmu = input_ids_mmu.to(accelerator.device, non_blocking=True)
image_tokens_mmu = vq_model.get_code(pixel_values_mmu)
image_tokens_mmu = image_tokens_mmu + len(uni_prompting.text_tokenizer)
input_ids_mmu = torch.cat([
(torch.ones(input_ids_mmu.shape[0], 1) * uni_prompting.sptids_dict['<|mmu|>']).to(
accelerator.device),
(torch.ones(input_ids_mmu.shape[0], 1) * uni_prompting.sptids_dict['<|soi|>']).to(
accelerator.device),
image_tokens_mmu,
(torch.ones(input_ids_mmu.shape[0], 1) * uni_prompting.sptids_dict['<|eoi|>']).to(
accelerator.device),
input_ids_mmu,
], dim=1).long()
labels_mmu = torch.cat([
(torch.ones(input_ids_mmu.shape[0], 1) * uni_prompting.ignore_id).to(accelerator.device),
(torch.ones(input_ids_mmu.shape[0], 1) * uni_prompting.ignore_id).to(accelerator.device),
torch.ones_like(image_tokens_mmu) * uni_prompting.ignore_id,
(torch.ones(input_ids_mmu.shape[0], 1) * uni_prompting.ignore_id).to(accelerator.device),
labels_mmu.to(accelerator.device)
], dim=1).long()
else:
pixel_values_mmu, texts_mmu = batch["mmu_flow"]["images"], batch["mmu_flow"]["input_ids"]
pixel_values_mmu = pixel_values_mmu.to(accelerator.device, non_blocking=True)
image_tokens_mmu = vq_model.get_code(pixel_values_mmu)
image_tokens_mmu = image_tokens_mmu + len(uni_prompting.text_tokenizer)
input_ids_mmu, prompt_masks, labels_mmu = uni_prompting((image_tokens_mmu, texts_mmu), 'mmu')
(
input_ids_mmu,
labels_mmu,
p_mask_mmu,
answer_lengths
) = prepare_inputs_and_labels_for_mmu(input_ids_mmu, prompt_masks, labels_mmu)
input_ids_mmu = input_ids_mmu.to(accelerator.device, non_blocking=True)
input_ids = torch.cat((input_ids, input_ids_mmu.to(input_ids.device)), dim=0)
labels = torch.cat((labels, labels_mmu.to(input_ids.device)), dim=0)
if global_step == 0 and epoch == 0:
logger.info("Input ids: {}".format(input_ids))
logger.info("Labels: {}".format(labels))
with accelerator.accumulate(model):
logits, loss_t2i, loss_lm, loss_mmu = model.forward_process(
input_ids=input_ids,
labels=labels,
batch_size_t2i=batch_size_t2i,
batch_size_lm=batch_size_lm,
batch_size_mmu=batch_size_mmu,
max_seq_length=config.dataset.preprocessing.max_seq_length,
p_mask_lm=p_mask_lm,
p_mask_mmu=p_mask_mmu,
answer_lengths=answer_lengths,
t2i_masks=t2i_masks
)
# Gather the losses across all processes for logging (if we use distributed training).
avg_loss_t2i = accelerator.gather(loss_t2i.repeat(config.training.batch_size_t2i)).mean()
avg_loss_lm = accelerator.gather(loss_lm.repeat(config.training.batch_size_lm)).mean()
avg_loss_mmu = accelerator.gather(loss_mmu.repeat(config.training.batch_size_mmu)).mean()
loss = config.training.t2i_coeff * loss_t2i + \
config.training.lm_coeff * loss_lm + \
config.training.mmu_coeff * loss_mmu
avg_masking_rate = accelerator.gather(mask_prob.repeat(config.training.batch_size_t2i)).mean()
accelerator.backward(loss)
if config.training.max_grad_norm is not None and accelerator.sync_gradients:
accelerator.clip_grad_norm_(model.parameters(), config.training.max_grad_norm)
optimizer.step()
lr_scheduler.step()
# log gradient norm before zeroing it
if (
accelerator.sync_gradients
and (global_step + 1) % config.experiment.log_grad_norm_every == 0
and accelerator.is_main_process
):
log_grad_norm(model, accelerator, global_step + 1)
optimizer.zero_grad(set_to_none=True)
# Checks if the accelerator has performed an optimization step behind the scenes
if accelerator.sync_gradients:
batch_time_m.update(time.time() - end)
end = time.time()
# Log metrics
if (global_step + 1) % config.experiment.log_every == 0:
samples_per_second_per_gpu = (
config.training.gradient_accumulation_steps * total_batch_size_per_gpu / batch_time_m.val
)
logs = {
"step_loss_t2i": avg_loss_t2i.item(),
"step_loss_mmu": avg_loss_mmu.item(),
"step_loss_lm": avg_loss_lm.item(),
"lr": lr_scheduler.get_last_lr()[0],
"avg_masking_rate": avg_masking_rate.item(),
"samples/sec/gpu": samples_per_second_per_gpu,
"data_time": data_time_m.val,
"batch_time": batch_time_m.val,
}
accelerator.log(logs, step=global_step + 1)
logger.info(
f"Step: {global_step + 1} "
f"Loss_t2i: {avg_loss_t2i.item():0.4f} "
f"Loss_mmu: {avg_loss_mmu.item():0.4f} "
f"Loss_lm: {avg_loss_lm.item():0.4f} "
f"Data (t): {data_time_m.val:0.4f}, {samples_per_second_per_gpu:0.2f}/s/gpu "
f"Batch (t): {batch_time_m.val:0.4f} "
f"LR: {lr_scheduler.get_last_lr()[0]:0.6f}"
)
# resetting batch / data time meters per log window
batch_time_m.reset()
data_time_m.reset()
# Save model checkpoint
if (global_step + 1) % config.experiment.save_every == 0:
save_checkpoint(model, config, accelerator, global_step + 1)
if ((global_step + 1) % config.experiment.generate_every == 0 or global_step == 0) and accelerator.is_main_process:
generate_images(
model,
vq_model,
uni_prompting,
accelerator,
config,
global_step + 1,
mask_schedule=mask_schedule,
)
visualize_predictions(
model,
vq_model,
uni_prompting,
config,
global_step + 1,
input_ids,
image_tokens_ori,
batch["t2i_flow"]["images"],
texts,
logits,
accelerator
)
understanding_images(
model,
vq_model,
uni_prompting,
accelerator,
config,
global_step + 1,
)
global_step += 1
if global_step >= config.training.max_train_steps:
break
accelerator.wait_for_everyone()
# Evaluate and save checkpoint at the end of training
save_checkpoint(model, config, accelerator, global_step)
# Save the final trained checkpoint
if accelerator.is_main_process:
model = accelerator.unwrap_model(model)
model.save_pretrained(config.experiment.output_dir, safe_serialization=True)
accelerator.end_training()
@torch.no_grad()
def visualize_predictions(
model,
vq_model,
uni_prompting,
config,
global_step,
input_ids,
image_tokens_ori,
ori_images,
texts,
logits,
accelerator
):
logger.info("Visualizing predictions...")
model.eval()
recons_images = vq_model.decode_code(image_tokens_ori - len(uni_prompting.text_tokenizer))
recons_images = torch.clamp((recons_images + 1.0) / 2.0, min=0.0, max=1.0)
recons_images *= 255.0
recons_images = recons_images.permute(0, 2, 3, 1).cpu().numpy().astype(np.uint8)
images = torch.clamp((ori_images + 1.0) / 2.0, min=0.0, max=1.0)
images *= 255.0
images = images.permute(0, 2, 3, 1).cpu().numpy().astype(np.uint8)
predictions = logits[:config.training.batch_size_t2i, -(config.model.mmada.num_vq_tokens + 1):-1:, len(uni_prompting.text_tokenizer) + config.model.mmada.num_new_special_tokens: len(uni_prompting.text_tokenizer) + config.model.mmada.num_new_special_tokens + config.model.mmada.codebook_size]
predictions = predictions.argmax(axis=-1)
mask_token_id = accelerator.unwrap_model(model).config.mask_token_id - len(uni_prompting.text_tokenizer)
input_ids = input_ids[:config.training.batch_size_t2i, -(config.model.mmada.num_vq_tokens + 1):-1:] - len(uni_prompting.text_tokenizer)
mask_ratio = list((torch.where(input_ids == mask_token_id, 1, 0).sum(
dim=-1) / config.model.mmada.num_vq_tokens).cpu().numpy())
predicted_images = torch.where(input_ids == mask_token_id, predictions, input_ids)
predicted_images = vq_model.decode_code(predicted_images)
predicted_images = torch.clamp((predicted_images + 1.0) / 2.0, min=0.0, max=1.0)
predicted_images *= 255.0
predicted_images = predicted_images.permute(0, 2, 3, 1).cpu().numpy().astype(np.uint8)
predicted_images = np.concatenate((images, recons_images, predicted_images), 2)
pil_images = [Image.fromarray(image) for image in predicted_images]
# Log images
wandb_images = [wandb.Image(image, caption=f'mask ratio: {r:0.2f} \n caption: {texts[i]}') for i, (image, r) in
enumerate(zip(pil_images, mask_ratio))]
wandb.log({"Original images v.s. Reconstructed images v.s. Predicted images": wandb_images}, step=global_step)
model.train()
@torch.no_grad()
def generate_images(
model,
vq_model,
uni_prompting,
accelerator,
config,
global_step,
mask_schedule,
):
logger.info("Generating images...")
model.eval()
# read validation prompts from file
with open(config.dataset.params.validation_prompts_file, "r") as f:
validation_prompts = f.read().splitlines()
mask_dtype = model.get_input_embeddings().weight.dtype
mask_token_id = accelerator.unwrap_model(model).config.mask_token_id
image_tokens = torch.ones((len(validation_prompts), config.model.mmada.num_vq_tokens), dtype=torch.long,
device=accelerator.device) * mask_token_id
input_ids, attention_mask = uni_prompting((validation_prompts, image_tokens), 't2i_gen')
if config.training.guidance_scale > 0:
uncond_input_ids, uncond_attention_mask = uni_prompting(([''] * len(validation_prompts), image_tokens), 't2i_gen')
else:
uncond_input_ids = None
uncond_attention_mask = None
if accelerator.mixed_precision == "fp16":
weight_dtype = torch.float16
elif accelerator.mixed_precision == "bf16":
weight_dtype = torch.bfloat16
else:
weight_dtype = torch.float32
with torch.autocast("cuda", dtype=weight_dtype, enabled=accelerator.mixed_precision != "no"):
# Generate images
gen_token_ids = accelerator.unwrap_model(model).t2i_generate(
input_ids=input_ids,
uncond_input_ids=uncond_input_ids,
attention_mask=attention_mask,
uncond_attention_mask=uncond_attention_mask,
guidance_scale=config.training.guidance_scale,
temperature=config.training.get("generation_temperature", 1.0),
timesteps=config.training.generation_timesteps,
noise_schedule=mask_schedule,
noise_type=config.training.get("noise_type", "mask"),
predict_all_tokens=config.training.get("predict_all_tokens", False),
seq_len=config.model.mmada.num_vq_tokens,
uni_prompting=uni_prompting,
config=config,
)
# In the beginning of training, the model is not fully trained and the generated token ids can be out of range
# so we clamp them to the correct range.
gen_token_ids = torch.clamp(gen_token_ids, max=accelerator.unwrap_model(model).config.codebook_size - 1, min=0)
images = vq_model.decode_code(gen_token_ids)
model.train()
if config.training.get("pre_encode", False):
del vq_model
# Convert to PIL images
images = torch.clamp((images + 1.0) / 2.0, min=0.0, max=1.0)
images *= 255.0
images = images.permute(0, 2, 3, 1).cpu().numpy().astype(np.uint8)
pil_images = [Image.fromarray(image) for image in images]
# Log images
wandb_images = [wandb.Image(image, caption=validation_prompts[i]) for i, image in enumerate(pil_images)]
wandb.log({"Generated images": wandb_images}, step=global_step)
@torch.no_grad()
def understanding_images(
model,
vq_model,
uni_prompting,
accelerator,
config,
global_step,
):
logger.info("Understanding images...")
model.eval()
file_list = os.listdir(config.dataset.params.mmu_image_root)
file_list = [f for f in file_list if f.lower().endswith(('.jpg', '.png', '.jpeg'))]
responses = ['' for i in range(len(file_list))]
images = []
device = accelerator.device
if accelerator.mixed_precision == "fp16":
weight_dtype = torch.float16
elif accelerator.mixed_precision == "bf16":
weight_dtype = torch.bfloat16
else:
weight_dtype = torch.float32
for i, file_name in enumerate(file_list):
image_path = os.path.join(config.dataset.params.mmu_image_root, file_name)
image_ori = Image.open(image_path).convert("RGB")
image = image_transform(image_ori, resolution=config.dataset.params.resolution).to(device)
image = image.unsqueeze(0)
images.append(image)
image_tokens = vq_model.get_code(image) + len(uni_prompting.text_tokenizer)
batch_size = 1
input_ids = uni_prompting.text_tokenizer(['<|start_header_id|>user<|end_header_id|>\n' + "Please describe this image in detail." +'<eot_id><|start_header_id|>assistant<|end_header_id|>\n'])['input_ids']
input_ids = torch.tensor(input_ids).to(device)
input_ids = torch.cat([
(torch.ones(input_ids.shape[0], 1) * uni_prompting.sptids_dict['<|mmu|>']).to(device),
(torch.ones(input_ids.shape[0], 1) * uni_prompting.sptids_dict['<|soi|>']).to(device),
image_tokens,
(torch.ones(input_ids.shape[0], 1) * uni_prompting.sptids_dict['<|eoi|>']).to(device),
(torch.ones(input_ids.shape[0], 1) * uni_prompting.sptids_dict['<|sot|>']).to(device),
input_ids
], dim=1).long()
with torch.autocast("cuda", dtype=weight_dtype, enabled=accelerator.mixed_precision != "no"):
output_ids = accelerator.unwrap_model(model).mmu_generate(input_ids)
# output_ids = torch.stack(output_ids).squeeze()[None]
text = uni_prompting.text_tokenizer.batch_decode(output_ids[:, input_ids.shape[1]:], skip_special_tokens=True)
responses[i] += text[0]
model.train()
images = torch.cat(images, dim=0)
images = torch.clamp((images + 1.0) / 2.0, min=0.0, max=1.0)
images *= 255.0
images = images.permute(0, 2, 3, 1).cpu().numpy().astype(np.uint8)
pil_images = [Image.fromarray(image) for image in images]
# Log images
wandb_images = [wandb.Image(image, caption=responses[i]) for i, image in enumerate(pil_images)]
wandb.log({"Understanding images": wandb_images}, step=global_step)
def save_checkpoint(model, config, accelerator, global_step):
output_dir = config.experiment.output_dir
checkpoints_total_limit = config.experiment.get("checkpoints_total_limit", None)
# _before_ saving state, check if this save would set us over the `checkpoints_total_limit`
if accelerator.is_main_process and checkpoints_total_limit is not None:
checkpoints = os.listdir(output_dir)
checkpoints = [d for d in checkpoints if d.startswith("checkpoint")]
checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1]))
# before we save the new checkpoint, we need to have at _most_ `checkpoints_total_limit - 1` checkpoints
if len(checkpoints) >= checkpoints_total_limit:
num_to_remove = len(checkpoints) - checkpoints_total_limit + 1
removing_checkpoints = checkpoints[0:num_to_remove]
logger.info(
f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints"
)
logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}")
for removing_checkpoint in removing_checkpoints:
removing_checkpoint = os.path.join(output_dir, removing_checkpoint)
shutil.rmtree(removing_checkpoint)
save_path = Path(output_dir) / f"checkpoint-{global_step}"
# retrieve the model on all processes for deepspeed stage 3 to work then save on one process (we are not using stage 3 yet)
# XXX: could also make this conditional on deepspeed
state_dict = accelerator.get_state_dict(model)
if accelerator.is_main_process:
unwrapped_model = accelerator.unwrap_model(model)
unwrapped_model.save_pretrained(
save_path / "unwrapped_model",
save_function=accelerator.save,
state_dict=state_dict,
safe_serialization=True
)
json.dump({"global_step": global_step}, (save_path / "metadata.json").open("w+"))
logger.info(f"Saved state to {save_path}")
def log_grad_norm(model, accelerator, global_step):
for name, param in model.named_parameters():
if param.grad is not None:
grads = param.grad.detach().data
grad_norm = (grads.norm(p=2) / grads.numel()).item()
accelerator.log({"grad_norm/" + name: grad_norm}, step=global_step)
if __name__ == "__main__":
main()