CogView

Generates images from Chinese text prompts by encoding both text and images as discrete token sequences and processing them through a unified 4-billion-parameter autoregressive transformer. The model treats image generation as a sequence prediction task, tokenizing images into 8192-code discrete tokens via a pretrained VQ-VAE, then autoregressively predicting image tokens conditioned on text token embeddings. This unified token-based approach enables the same model weights to support multiple downstream tasks (generation, captioning, super-resolution) without task-specific architectures.

Upscales low-resolution images by tokenizing them with the same VQ-VAE encoder, then using the cogview-sr checkpoint to autoregressively predict higher-resolution token sequences. The model learns to map low-res token distributions to high-res token distributions within the discrete token space, preserving semantic content while increasing visual fidelity. This approach avoids pixel-space upsampling artifacts by operating entirely in the learned token manifold.

Implements efficient batch inference via generate_samples.py with dynamic batch size adjustment based on available GPU memory. The inference pipeline accepts --max-inference-batch-size parameter, which is automatically reduced if GPU memory is insufficient, enabling inference on GPUs with less than V100 VRAM. Batching is implemented via PyTorch's DataLoader with custom collation, enabling efficient processing of multiple prompts/images in parallel.

Provides evaluation utilities (in utils.py) for computing metrics on generated images, including image quality scores (via pretrained perceptual models) and text-image alignment scores (via the cogview-caption model). These utilities enable quantitative evaluation of generation quality without human review, supporting both single-image and batch evaluation modes. Metrics are computed in discrete token space when possible, avoiding pixel-space artifacts.

Generates natural language captions for images by tokenizing them with the VQ-VAE encoder, then using the cogview-caption checkpoint to autoregressively predict Chinese text tokens conditioned on image tokens. The model learns bidirectional image-to-text mapping within the unified token space, enabling the same transformer weights to generate descriptive captions from visual input. This reverses the text-to-image direction while maintaining the same autoregressive decoding mechanism.

Scores and ranks multiple generated images using the cogview-caption checkpoint as a preference model, computing relevance scores between image tokens and the original text prompt. The model encodes both the image and text as token sequences, then uses transformer attention to compute alignment scores that reflect how well each image matches the input prompt. This enables selection of the best image from a batch of candidates without additional model inference.

Stabilizes large-scale transformer training by mitigating floating-point overflow in attention computation during mixed-precision (FP16/FP32) training. PB-relax dynamically adjusts the precision of attention logits to prevent overflow while maintaining gradient flow, implemented via custom CUDA kernels in the attention module. This technique is configured in arguments.py and active by default in pretrained checkpoints, enabling stable training of 4B-parameter models without NaN losses.

Stabilizes deep transformer training by placing layer normalization in a sandwich pattern (pre-norm and post-norm) rather than standard pre-norm or post-norm alone. This alternative normalization placement eliminates NaN losses and improves gradient flow in deep networks, implemented as a configurable layer norm variant in the transformer blocks. Sandwich-LN is active by default in pretrained checkpoints and is configured via arguments.py, enabling training of very deep transformers without numerical instability.

Enables training of 4B-parameter models across multiple GPU nodes using DeepSpeed's ZeRO (Zero Redundancy Optimizer) stage 2/3, which partitions model parameters, gradients, and optimizer states across devices to reduce per-GPU memory usage. The training pipeline integrates DeepSpeed's distributed communication primitives (AllReduce, AllGather) with PyTorch's DistributedDataParallel, configured via arguments.py with node count, rank, and backend settings. This enables scaling to multi-node clusters while maintaining convergence.

Preprocesses training data by encoding images into discrete token sequences using a pretrained VQ-VAE (vqvae_hard_biggerset_011.pt), which maps images to 8192-code tokens via learned quantization. The data pipeline (implemented in data_utils.py and dataset classes) handles both image tokenization and text tokenization (via SentencePiece), creating aligned token sequences for transformer training. This enables efficient batching and caching of tokenized data, reducing per-epoch preprocessing overhead.

Centralizes all training, inference, and model configuration in arguments.py, which defines command-line arguments for model architecture (depth, width, attention type), training hyperparameters (learning rate, batch size, warmup), distributed settings (node rank, world size), and stability techniques (PB-relax, Sandwich-LN). The argument parser is used by all entry points (generate_samples.py for inference, training scripts for training), enabling reproducible configuration management and easy hyperparameter sweeps via command-line overrides.

Implements checkpoint saving and loading that handles distributed training state, including model parameters, optimizer state, and training metadata (epoch, step, loss). The checkpointing system (in utils.py) ensures that all distributed ranks save/load synchronized state, preventing data corruption from asynchronous writes. Checkpoints include model architecture configuration, enabling resumption of training from arbitrary steps with full state recovery.