import os import shutil import signal import sys import time from contextlib import ExitStack from functools import partial from pathlib import Path from accelerate.utils import gather_object, gather from torchinfo import summary from unidisc.tokenizers.chameleon_tokenizers import tokenize_chameleon, tokenize_chameleon_fast, get_chameleon_images, decode_ids, decode_ids_batched, tokenize_chameleon_mmc4, tokenize_regular_cambrian_mmc4 from utils import _print_config, set_numa_affinity, set_omega_conf_resolvers sys.path.append(str(Path(__file__).parent.parent.parent / "unidisc/misc/hydra_submitit_launcher")) import json import os import random import sys from contextlib import nullcontext from pathlib import Path import fsspec import hydra import numpy as np import omegaconf import rich.syntax import rich.tree import torch from accelerate import Accelerator from PIL import Image from tensordict import TensorDict from tqdm import tqdm try: from viztracer import VizTracer except ImportError: print("VizTracer not installed, skipping") from dataloader import get_dataloaders, get_tokenizer, tokenize_text from decoupled_utils import (barrier, breakpoint_on_error, get_local_rank, get_rank, get_world_size, is_local_main_process, is_main_process, rank_zero_fn, rprint, set_global_breakpoint, set_global_exists, gprint) from model import decode_latents, get_image_batch, get_vae from models.datasets.combine_token_dicts import main as combine_token_dicts from models.datasets.vggface_v2_attributes import (get_inference_func, get_output) from utils import (_print_config, set_numa_affinity, set_omega_conf_resolvers, set_torch_defaults) from omegaconf import DictConfig, OmegaConf, open_dict, read_write os.environ["HYDRA_FULL_ERROR"] = "1" set_global_breakpoint() # Overrides breakpoint() to use ipdb.set_trace() instead and handle distributed training set_global_exists() set_omega_conf_resolvers() set_torch_defaults() def get_batch_size(config): with open_dict(config): if any(x.lower() in torch.cuda.get_device_name().lower() for x in ["v100", "1080", "2080", "quadro", "titan"]) or torch.cuda.get_device_capability()[0] <= 7: config.trainer.precision = "no" config.model.force_optimized_native_attn = False config.trainer.compile = False config.loader.batch_size = config.loader.batch_size // 3 print(f"Found {torch.cuda.get_device_name().lower()}, set batch size to {config.loader.batch_size}") return config def enc(data, idx, encode_images, config, vae, batch, accelerator, mixed_precision, tokenizer, vgg_data, existing_ids=None, device=None, mapping=None): if isinstance(batch, list): bs = len(batch) elif "img" in batch: bs = batch["img"].shape[0] else: bs = batch["attention_mask"].shape[0] sl = slice(idx * bs, (idx + 1) * bs) if not isinstance(batch, list) and "idx" in batch: if set(data[sl]["idx"].flatten().tolist()) == set(batch["idx"].tolist()): rprint(f"Skipping {idx} as all samples have already been processed 1") return if existing_ids is not None: set_inter = set(batch["idx"].tolist()) & existing_ids if len(set_inter) == bs: rprint(f"Skipping {idx} as all samples have already been processed 2") return elif len(set_inter) > 0: rprint(f"Running {idx} as some samples have already been processed: {len(set_inter)}") else: if (data[sl]["idx"] != -1).all(): rprint(f"Skipping {idx} as all samples have already been processed") return if not isinstance(batch, list) and "img" in batch: batch["img"] = batch["img"].to(device=device, dtype=torch.bfloat16 if mixed_precision else None) with torch.no_grad(): with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=mixed_precision): use_chameleon = getattr(config.data, "use_chameleon", False) use_mmc4 = config.data.train == "mmc4" use_cambrian = config.data.train == "cambrian" if not use_chameleon and not use_mmc4 and not use_cambrian: if tokenizer is not None and getattr(config.model, "unified_model", False): if "input_ids" in batch and "attention_mask" in batch: tokens = batch else: tokens = tokenize_text(tokenizer, config.data.block_size, batch[".txt"]) batch["txt_input_ids"] = tokens["input_ids"] batch["txt_attention_mask"] = tokens["attention_mask"].float() elif getattr(config.data, "add_vggface_v2_attributes", False) and "vggface" not in config.data.train: txt_input_ids, txt_attention_mask = get_output(batch, **vgg_data) batch["txt_input_ids"] = txt_input_ids batch["txt_attention_mask"] = txt_attention_mask elif getattr(config.data, "txt_only", False): batch["txt_input_ids"] = batch["input_ids"] batch["txt_attention_mask"] = batch["attention_mask"] if getattr(config.model, "unified_model", False) is False: if getattr(config.data, "txt_only", False): batch["modality"] = torch.full((bs, 1), fill_value=0, dtype=torch.int16) else: batch["modality"] = torch.full((bs, 1), fill_value=1, dtype=torch.int16) if isinstance(batch, list) and batch[0].get("idx", None) is not None: _idx = torch.tensor([x["idx"] for x in batch], dtype=torch.int32).unsqueeze(-1) elif "idx" in batch: _idx = batch["idx"].to(torch.int32).unsqueeze(-1) else: _idx = torch.full((bs, 1), fill_value=0, dtype=torch.int32) if "is_valid" in batch: _idx[~batch["is_valid"]] = -1 if (_idx == -1).all(): gprint(f"WARNING: All samples are invalid") sl = slice(idx * bs, (idx + 1) * bs) assert (idx + 1) * bs <= len(data), f"Index {idx} + batch size {bs} is greater than the data length {len(data)}" if encode_images: if use_chameleon: if isinstance(batch, list): all_input_ids, all_attention_masks = tokenize_chameleon_mmc4(config, tokenizer, vae, batch, device, mapping) else: all_input_ids, all_attention_masks = tokenize_chameleon_fast(config, tokenizer, vae, batch) # all_input_ids_gt, all_attention_masks_gt = tokenize_chameleon(config, tokenizer, vae, batch) # txt_tokens, img_tokens = decode_ids_batched(_vae, all_input_ids[:4], return_tokens=True) # img = decode_latents(config, _vae, img_tokens) # from image_utils import Im; Im(img).save() elif use_mmc4 or use_cambrian: all_input_ids, all_attention_masks, all_modality = tokenize_regular_cambrian_mmc4(config, tokenizer, vae, batch, device, mapping) if all_input_ids is None: return else: image_ids = get_image_batch(config, vae, batch, device) if use_chameleon or use_mmc4 or use_cambrian: if not use_chameleon: assert (all_input_ids < torch.iinfo(torch.int16).max).all() _kwargs = {} if use_mmc4 or use_cambrian: _kwargs["modality"] = all_modality.to(torch.int8) data[sl] = TensorDict( { "input_ids": all_input_ids.to(torch.int32 if use_chameleon else torch.int16), "attention_mask": all_attention_masks.to(torch.bool), "idx": _idx, "write_flag": torch.ones((bs, 1), dtype=torch.bool), **_kwargs, }, batch_size=[bs], ) elif getattr(config.model, "cond_label", False): data[sl] = TensorDict( { "img_input_ids": image_ids.to(torch.int16), "img_label": batch["label"].to(torch.int32).unsqueeze(-1), "idx": _idx, "write_flag": torch.ones((bs, 1), dtype=torch.bool), }, batch_size=[bs], ) elif getattr(config.model, "unified_model", False) or getattr(config.data, "add_vggface_v2_attributes", False): data[sl] = TensorDict( { "img_input_ids": image_ids.to(torch.int16), "txt_input_ids": (batch.get("txt_input_ids") if batch.get("txt_input_ids") is not None else batch["input_ids"]).to( torch.int32 ), "txt_attention_mask": ( batch.get("txt_attention_mask") if batch.get("txt_attention_mask") is not None else batch["attention_mask"] ).to(torch.bool), "idx": _idx, "write_flag": torch.ones((bs, 1), dtype=torch.bool), }, batch_size=[bs], ) else: data[sl] = TensorDict( {"input_ids": image_ids.to(torch.int32), "attention_mask": torch.ones((image_ids.shape[0], image_ids.shape[1]), dtype=torch.bool), "idx": _idx, "write_flag": torch.ones((bs, 1), dtype=torch.bool), "modality": batch["modality"].to(torch.int16)}, batch_size=[bs], ) elif getattr(config.data, "txt_only", False): data[sl] = TensorDict( {"input_ids": batch['input_ids'].to(torch.int32), "attention_mask": batch['attention_mask'].to(torch.bool), "idx": _idx, "write_flag": torch.ones((bs, 1), dtype=torch.bool), "modality": batch["modality"].to(torch.int16)}, batch_size=[bs], ) else: real_image = batch["img"] if (config.data.resolution == 512 and batch["img"].shape[0] > 16) or (config.model.downscale_ratio <= 8): chunk_size = 8 if (config.model.image_vocab_size > 64000 or config.model.downscale_ratio <= 8) else 16 chunks = [batch["img"][i : i + chunk_size] for i in range(0, batch["img"].shape[0], chunk_size)] rec_img_list = [] for chunk in chunks: batch_chunk = {"img": chunk} image_ids = get_image_batch(config, vae, batch_chunk, device) rec_img = decode_latents(config, vae, image_ids) rec_img_list.append(rec_img) rec_img = torch.cat(rec_img_list, dim=0) else: image_ids = get_image_batch(config, vae, batch, device) rec_img = decode_latents(config, vae, image_ids) viz_img = torch.cat([real_image, rec_img], dim=-1) from image_utils import Im if getattr(config.model, 'custom_vae_name', None) is not None: custom_str = getattr(config.model, 'custom_vae_name') else: custom_str = f"{'_custom' if getattr(config.model, 'use_custom_vae_ckpt', False) else ''}" (Path(__file__).parent.parent.parent / "output").mkdir(parents=True, exist_ok=True) Im(viz_img).save( Path(__file__).parent.parent.parent / f"output/{config.data.train.replace('/', '')}_seq{image_ids.shape[1]}_res{config.data.resolution}_{config.model.vae_type}{custom_str}_voc{config.model.image_vocab_size}.png" ) # Create directories for saving images dataset_name = config.data.train.replace('/', '') vae_name = f"seq{image_ids.shape[1]}_res{config.data.resolution}_{config.model.vae_type}{custom_str}_voc{config.model.image_vocab_size}" output_dir = Path(__file__).parent.parent.parent / "output" / dataset_name / vae_name gt_output_dir = Path(__file__).parent.parent.parent / "output" / dataset_name / f"GT_{config.data.resolution}" output_dir.mkdir(parents=True, exist_ok=True) gt_output_dir.mkdir(parents=True, exist_ok=True) # Save each image separately for i, (real, rec) in enumerate(zip(real_image, rec_img)): print(Im(rec).save(output_dir / f"{i}.png")) if (gt_output_dir / f"{i}.png").exists() is False: print(Im(real).save(gt_output_dir / f"{i}.png")) gprint(f"Exiting") exit() def get_dict(config, dataset_size): if getattr(config.data, "use_chameleon", False) or config.data.train == "cambrian" or config.data.train == "mmc4": input_ids_dtype = torch.int32 if getattr(config.data, "use_chameleon", False) else torch.int16 data = TensorDict( { "input_ids": torch.zeros(dataset_size, config.model.length, dtype=input_ids_dtype), "attention_mask": torch.zeros(dataset_size, config.model.length, dtype=torch.bool), "modality": torch.full((dataset_size, config.model.length), fill_value=-1, dtype=torch.int8), "idx": torch.full((dataset_size, 1), fill_value=-1, dtype=torch.int32), "write_flag": torch.zeros(dataset_size, 1, dtype=torch.bool), }, batch_size=[dataset_size], ) elif getattr(config.model, "cond_label", False): data = TensorDict( { "img_input_ids": torch.zeros(dataset_size, config.model.img_length, dtype=torch.int16), "img_label": torch.zeros(dataset_size, 1, dtype=torch.int32), "idx": torch.full((dataset_size,), fill_value=-1, dtype=torch.int32), "write_flag": torch.zeros(dataset_size, 1, dtype=torch.bool), }, batch_size=[dataset_size], ) elif getattr(config.model, "unified_model", False) or getattr(config.data, "add_vggface_v2_attributes", False): data = TensorDict( { "img_input_ids": torch.zeros(dataset_size, config.model.img_length, dtype=torch.int16), "txt_input_ids": torch.zeros(dataset_size, config.model.txt_length, dtype=torch.int32), "txt_attention_mask": torch.zeros(dataset_size, config.model.txt_length, dtype=torch.bool), "idx": torch.full((dataset_size, 1), fill_value=-1, dtype=torch.int32), "write_flag": torch.zeros(dataset_size, 1, dtype=torch.bool), }, batch_size=[dataset_size], ) else: data = TensorDict( { "input_ids": torch.zeros(dataset_size, config.model.txt_length if config.data.txt_only else config.model.img_length, dtype=torch.int16), "idx": torch.full((dataset_size, 1), fill_value=-1, dtype=torch.int32), "write_flag": torch.zeros(dataset_size, 1, dtype=torch.bool), "modality": torch.full((dataset_size, 1), fill_value=-1, dtype=torch.int16), }, batch_size=[dataset_size], ) return data def signal_handler(signum, frame, train_data, tmp_path): """Handle signals to save temporary train data.""" rprint(f"Received signal {signum}, saving temporary train data.") print(f"[PRINT] Received signal {signum}, saving temporary train data.") save_tmp_data(train_data, tmp_path) sys.exit def save_tmp_data(data, tmp_path): """Save data to a temporary path.""" if tmp_path.exists() and tmp_path.is_dir(): rprint(f"Deleting {tmp_path}") shutil.rmtree(tmp_path) # Delete old tmp directory if it exists rprint(f"Saving tmp data to {tmp_path}") data.memmap(tmp_path, copy_existing=True) def periodic_save(data, tmp_path, start_time, interval=2 * 60 * 60): """Periodically save data to a temporary path.""" current_time = time.time() if current_time - start_time >= interval: rprint(f"Hit periodic save interval, saving tmp data to {tmp_path}") save_tmp_data(data, tmp_path) return current_time # Reset start time return start_time @hydra.main(version_base=None, config_path="../../configs", config_name="config") def main(config): """Main entry point for training.""" try: import resource soft_limit, hard_limit = resource.getrlimit(resource.RLIMIT_NOFILE) resource.setrlimit(resource.RLIMIT_NOFILE, (hard_limit, hard_limit)) # Set the soft limit to the hard limit rprint(f"Successfully set RLIMIT_NOFILE to {hard_limit}") except Exception as e: rprint(f"Failed to set RLIMIT_NOFILE: {e}") mixed_precision = False train_start_time = time.time() from datetime import timedelta from accelerate import Accelerator, DataLoaderConfiguration from accelerate.utils import InitProcessGroupKwargs kwargs = InitProcessGroupKwargs(timeout=timedelta(seconds=3600)) prepare_kwargs = {} if config.data.train == "mmc4": prepare_kwargs["dispatch_batches"] = False accelerator = Accelerator(mixed_precision="bf16" if mixed_precision else None, kwargs_handlers=[kwargs], dataloader_config=DataLoaderConfiguration(**prepare_kwargs)) device = torch.device(f"cuda:{accelerator.local_process_index}") import socket hostname = socket.gethostname() print(f"Hostname: {hostname}, Process index: {accelerator.process_index}, {device}, local_process_index: {accelerator.local_process_index}, get_local_process_index: {get_local_rank()}, device: {device}") _print_config(config, resolve=True, save_cfg=True) config = get_batch_size(config) # with omegaconf.open_dict(config): # batch_sizes = gather_object([config.loader.batch_size]) # rprint(f"Batch sizes: {batch_sizes}") # smallest_batch_size = min(batch_sizes) # config.loader.batch_size = smallest_batch_size # rprint(f"New config batch size: {config.loader.batch_size}") prefix = f"[Rank {accelerator.process_index}/{accelerator.num_processes}, Node: {os.environ.get('SLURM_NODEID', 'N/A')}, Hostname: {os.environ.get('SLURM_JOB_NODELIST', 'N/A')}, {config.data.train}]" print(f"{prefix} Starting precomputing tokens") save_validation_dataloader = getattr(config.data, "save_validation_dataloader", False) save_train_dataloader = getattr(config.data, "save_train_dataloader", False) tokenizer = get_tokenizer(config) train_dataloader, val_dataloader = get_dataloaders( config, tokenizer=tokenizer, allow_aug=False, force_aug=getattr(config.data, "force_aug", False), skip_valid=not save_validation_dataloader ) train_dataloader = accelerator.prepare(train_dataloader) if save_validation_dataloader: val_dataloader = accelerator.prepare(val_dataloader) encode_images = getattr(config.model, "encode_images", False) use_chameleon = getattr(config.data, "use_chameleon", False) use_mmc4 = config.data.train == "mmc4" use_cambrian = config.data.train == "cambrian" mapping = None if use_chameleon: from unidisc.tokenizers.chameleon_tokenizers import ItemProcessor vae = ItemProcessor(target_size=config.data.resolution) else: vae = get_vae(config, device) if use_mmc4: import pandas as pd mapping = pd.read_parquet(config.data.mmc4_mapping_parquet) # Keep tar_filepath if it exists, otherwise use shard_path or map img2dataset_shard_id if "tar_filepath" in mapping.columns: pass elif "shard_path" in mapping.columns: mapping = mapping.rename(columns={"shard_path": "tar_filepath"}) mapping["tar_filepath"] = mapping["tar_filepath"].str.replace(".parquet", ".tar") else: tar_path = Path(config.data.mmc4_tar_path) mapping["tar_filepath"] = mapping["img2dataset_shard_id"].apply(lambda x: tar_path / f"{x}.tar") mapping = mapping[['url', 'tar_filepath', 'key']] mapping = mapping.set_index("url").sort_index() if use_mmc4 or use_cambrian: assert config.data.use_slow_tokenizer and config.data.add_image_token if config.data.iterable: train_dataset_size = getattr(config.data, "train_dataset_size", None) else: print(f"{prefix} Train dataloader: {len(train_dataloader)} batches") print(f"{prefix} Train underlying dataset: {len(train_dataloader.dataset)} samples") train_dataset_size = (len(train_dataloader.dataset) // accelerator.num_processes) + config.loader.batch_size if save_validation_dataloader: print(f"{prefix} Val dataloader: {len(val_dataloader)} batches") print(f"Val underlying dataset: {len(val_dataloader.dataset)} samples") val_dataset_size = (len(val_dataloader.dataset) // accelerator.num_processes) + config.loader.batch_size print(f"{prefix} Train dataset size: {train_dataset_size} for 1 GPU") if save_validation_dataloader: print(f"{prefix} Val dataset size: {val_dataset_size} for 1 GPU") rank = accelerator.process_index output_dir = config.data.token_output_dir output_dir = Path(f"{output_dir}") output_dir.mkdir(parents=True, exist_ok=True) assert config.data.force_disable_shuffle debug = getattr(config.data, "debug", False) print(f"{prefix} Output dir: {output_dir}") vgg_data = None if getattr(config.data, "add_vggface_v2_attributes", False): print(f"{prefix} Adding VGGFace V2 attributes") vgg_data = get_inference_func() vgg_data["model"] = accelerator.prepare(vgg_data["model"]) if not config.data.split_dataset and is_main_process() and any(output_dir.iterdir()): rprint(f"Found temporary directories in output dir, combining them") combine_token_dicts(output_dir, use_tmp=False, use_timestamp=True, delete_after_combining=True) for item in output_dir.iterdir(): if item.is_dir() and "tmp" in item.name: rprint(f"Removing temporary directory: {item}") shutil.rmtree(item) # barrier() # TODO: Should be a barrier here if not config.data.split_dataset: existing_folders = sorted([folder for folder in output_dir.iterdir() if folder.is_dir() and "existing" in folder.name]) if existing_folders: rprint(f"Found existing folders: {existing_folders}") existing_data = torch.cat([TensorDict.load_memmap(folder) for folder in existing_folders], dim=0) rprint(f"Concatenated existing data with shape: {existing_data.shape}") existing_ids = set(existing_data["idx"].to(torch.int32).flatten().tolist()) else: rprint("No existing folders found") existing_ids = None else: existing_ids = None if save_train_dataloader: if not config.data.split_dataset and getattr(config.data, "allow_load_from_tmp", True) and Path(output_dir / f"tmp_train_{rank}").exists(): rprint("Found tmp_train_{rank} in output dir, loading from it") train_data = TensorDict.load_memmap(output_dir / f"tmp_train_{rank}") train_data = train_data.clone() else: train_data = get_dict(config, train_dataset_size) print(f"{prefix} Starting train dataloader") if config.data.split_dataset: rank = int(os.getenv("SLURM_ARRAY_TASK_ID")) print(f"Using task id: {rank}") split_path = output_dir / f"train_{rank}" tmp_train_path = output_dir / f"tmp_train_{rank}" signal.signal(signal.SIGUSR1, partial(signal_handler, train_data=train_data, tmp_path=tmp_train_path)) signal.signal(signal.SIGUSR2, partial(signal_handler, train_data=train_data, tmp_path=tmp_train_path)) try: signal.signal(signal.SIGKILL, partial(signal_handler, train_data=train_data, tmp_path=tmp_train_path)) except: rprint(f"Failed to set SIGKILL handler") start_time = time.time() with VizTracer(output_file="optional.json", tracer_entries=5000000) if debug else nullcontext(): for i, batch in tqdm(enumerate(train_dataloader), leave=False, disable=not is_local_main_process()): if i == 0 and "img" in batch: print(f"Batch shape: {batch['img'].shape}") if debug and i >= 1: break enc(train_data, i, encode_images, config, vae, batch, accelerator, mixed_precision, tokenizer, vgg_data=vgg_data, existing_ids=existing_ids, device=device, mapping=mapping) try: if not config.data.split_dataset or True: start_time = periodic_save(train_data, tmp_train_path, start_time, getattr(config.data, "periodic_save", 2 * 60 * 60)) except Exception as e: gprint(f"Failed to save train data: {e}") start_time = time.time() if debug: exit() del train_dataloader print(f"{prefix} Saving train data") if split_path.exists() and split_path.is_dir(): rprint(f"Removing {split_path}") shutil.rmtree(split_path) split_path.mkdir(parents=True, exist_ok=True) gprint(f"Saving train data to {split_path}: {train_data.shape}") train_data.memmap(split_path, copy_existing=True) if tmp_train_path.exists() and tmp_train_path.is_dir(): rprint(f"Removing {tmp_train_path}") shutil.rmtree(tmp_train_path) if not config.data.split_dataset: with open(output_dir / f"train_{rank}.completed", 'w') as f: f.write(f"Processing done for rank {rank}\n") print(f"{prefix} Finished train dataloader") if save_validation_dataloader: val_data = get_dict(config, val_dataset_size) split_path = output_dir / f"val_{rank}" split_path.mkdir(parents=True, exist_ok=True) tmp_val_path = output_dir / f"tmp_val_{rank}" print(f"Starting val dataloader") start_time = time.time() # Track start time for periodic saving for i, batch in tqdm(enumerate(val_dataloader), leave=False): if debug and i >= 10: break enc(val_data, i, encode_images, config, vae, batch, accelerator, mixed_precision, tokenizer, vgg_data=vgg_data, device=device) # Periodically save data start_time = periodic_save(val_data, tmp_val_path, start_time) print(f"{prefix} Saving val data") if split_path.exists() and split_path.is_dir(): rprint(f"Removing {split_path}") shutil.rmtree(split_path) split_path.mkdir(parents=True, exist_ok=True) rprint(f"Saving val data to {split_path}") val_data.memmap(split_path, copy_existing=True) if tmp_val_path.exists() and tmp_val_path.is_dir(): shutil.rmtree(tmp_val_path) # Delete tmp directory after final save print(f"{prefix} Finished val dataloader") rprint(f"{prefix} Finished precomputing tokens") if config.data.split_dataset: rprint(f"We are splitting the dataset and thus exiting.") exit() if get_world_size() > 1 and (time.time() - train_start_time) > 60 * 60: time.sleep(60 * 60) barrier() rprint('after barrier') if is_main_process(): combine_token_dicts(data_dir=output_dir, allow_zero_idx=True, move_files=True, delete_after_combining=True) barrier() rprint(f"Finished concating tokens") if __name__ == "__main__": with breakpoint_on_error(): main()