lxiangyue/rosetta

Rosetta is a research project and codebase from HKUST and Tencent Hunyuan for training AI models that can handle text, images, and image generation together, rather than bolting these abilities on separately. Many models that try to add new abilities, like generating images, on top of an existing language model end up forgetting some of what they already knew about language, a problem the authors call catastrophic forgetting. Rosetta is built to avoid this by sharing a core of knowledge across all the tasks, while still letting each type of task, text, image understanding, and image generation, use specialized parts of the network suited to it. The paper behind this repository describes an architecture called Rosetta FFN with three main ideas. Attention layers are shared globally so information can flow between different types of input. Feed forward layers are split into specialized experts for text, images, and image generation, but all of them connect back to one shared expert that anchors the model's foundational knowledge. Finally, an optimization technique described in the paper resolves conflicting training signals between tasks without needing extra memory. Using this repository requires a serious machine learning setup. It expects Python 3.12 and a recent CUDA toolkit for GPU acceleration, plus PyTorch, Flash Attention, and several gigabytes of downloaded model weights, tokenizers, and evaluation datasets from Hugging Face. The authors provide a one command demo that reproduces the forgetting problem directly, training a model to generate images for just 100 steps and then comparing how much its language ability drops compared to two competing designs, called MoE and MoT. According to the README, Rosetta holds onto its language ability far better than either alternative in this test. The repository also includes scripts for full training runs, both on a single GPU and across multiple machines with multiple GPUs each, along with instructions for running evaluation benchmarks like MMLU and ARC Challenge. This is research code aimed at people already working on large multimodal models rather than a beginner friendly tool, and it is released under the Apache 2.0 license.