acytoo

LLM Tutorial 3 - Operator Acceleration and Native Multimodality: Computing Chinese, English, and Piano Music Simultaneously with 12.5M Parameters

#Operator Acceleration and Native Multimodality: Computing Chinese, English, and Piano Music Simultaneously with 12.5M Parameters

Series Introduction: In the first two posts (Part 1: From Bigram to GPT-2 and Part 2: Computing Piano Music with GPT-2), we built GPT-2 from scratch with straightforward, readable code and used it to generate Shakespearean text and MIDI piano music separately. However, that pedagogical implementation broke tensor operations into too many tiny steps, resulting in low GPU efficiency. In this post, we first switch to Andrej Karpathy's standard nanoGPT implementation, resolving operator bottlenecks by fusing QKV projections and introducing FlashAttention. Then, without changing a single line of network architecture, we pack Chinese, English, and piano music into the exact same vocabulary to train a native multimodal model with just 12.51 million (12.51M) parameters, running cross-modal "Text-to-Music" and "Music-to-Text" generation end-to-end.

#1. Operator Acceleration: Why Was Our Pedagogical Code Slow?

In SimpleLLM from Part 1, to keep a clean 1-to-1 correspondence between the code and the math equations, we split Self-Attention into separate linear layers per head and explicit matrix multiplications. While easy to follow, on modern GPUs this hits the Memory Wall (VRAM bandwidth bottleneck) and incurs heavy CUDA kernel launch overhead.

Compared to our pedagogical code, Karpathy's nanoGPT produces mathematically identical results while making three core low-level refactorings:

  1. Fused QKV Projection (c_attn): Instead of using three separate nn.Linear(d, d) layers to compute $Q, K, V$ individually, they are merged into a single nn.Linear(d, 3 * d). After one large matrix multiplication, .split(n_embd, dim=2) slices out $Q, K, V$, which are reshaped directly into 4D tensors of shape (B, nh, T, hs) to compute all attention heads in parallel. This cuts small-matrix kernel launch overhead by $2/3$ and avoids reading the input tensor multiple times.
  2. Replacing Manual Attention with FlashAttention: In manual attention (q @ k.T $\to$ masked_fill $\to$ softmax $\to$ @ v), every step writes an intermediate matrix of shape (B, nh, T, T) back to slow GPU High Bandwidth Memory (HBM) and reads it back in the next step. Switching to PyTorch's native F.scaled_dot_product_attention(..., is_causal=True) uses tiled computation to evaluate causal attention in a single pass inside the GPU's ultra-fast on-chip SRAM, never materializing the $T \times T$ attention matrix in HBM.
  3. Vocabulary Padding and Weight Tying: The input token embedding table wte and the output projection lm_head point to the exact same weight tensor in memory. Meanwhile, the vocabulary size is padded up to the nearest multiple of 64 (e.g., padding GPT-2's 50257 to 50304) to align with GPU Tensor Core hardware tile boundaries.

In the first Colab notebook, we load OpenAI's official gpt2-large (774M) pretrained weights from HuggingFace directly into this <300-line implementation and run text continuation on a free T4 GPU (FP16), verifying exact architectural alignment.

#2. How Do We Pack Chinese, English, and Piano Music into a Single Model?

Part 2 already proved that to a Transformer, both text characters and piano keys are just 1D integer IDs (Token IDs). So if we want a single model to learn Chinese, English, and music simultaneously—and even translate between them—do we need to design a specialized "music encoder" or "language decoder"?

Not at all. We simply pack the discrete symbols of all modalities into a single shared vocabulary (vocab_size = 4,928) and use a few special control tokens to mark modality transitions:

  • 0 .. 307 (MIDI Piano Events, 308 tokens): Reuses the exact encoding from Part 2: 0..87 for key press NOTE_ON, 88..175 for key release NOTE_OFF, 176..275 for TIME_SHIFT (10 ms per step), and 276..307 for 32-level key strike VELOCITY.
  • 308 .. 311 (Modality Control Tokens, 4 tokens):
    • 308 (TOK_TEXT): Marks the start of the text modality;
    • 309 (TOK_MUSIC): Marks the start of the music modality;
    • 310 (TOK_EOS): End-of-sequence token;
    • 311 (TOK_UNK): Unknown character placeholder.
  • 312 .. 4887 (Chinese & English Characters, 4,576 tokens): Indexes all unique characters appearing in the training corpora (the classic Chinese novel Journey to the West [《西游记》], the English novel Little Brother, and bilingual music descriptions).
  • 4888 .. 4927 (Alignment Padding, 40 tokens): Pads the raw vocabulary size 4,888 up to the nearest multiple of 64: 4,928.

With this unified encoding, both unimodal continuation and cross-modal generation become the exact same 1D next-integer prediction game to the model:

  • Unimodal Text Continuation: [308(TOK_TEXT), Chinese/English character sequence...]
  • Unimodal Piano Generation: [309(TOK_MUSIC), MIDI event sequence..., 310(TOK_EOS)]
  • Text-to-Music ($\text{Text} \to \text{Music}$): [308(TOK_TEXT), music description characters..., 309(TOK_MUSIC), MIDI event sequence..., 310(TOK_EOS)]
  • Music-to-Text ($\text{Music} \to \text{Text}$): [308(TOK_TEXT), "Describe this music in Chinese: ", 309(TOK_MUSIC), MIDI event sequence..., 308(TOK_TEXT), music description characters..., 310(TOK_EOS)]

In cross-modal tasks, to prevent a long prefix (such as 200–300 MIDI note tokens) from diluting the gradients of the target output, when constructing the target labels y we set all prompt prefix tokens before the modality delimiter to -1 (leveraging ignore_index = -1 in F.cross_entropy to skip loss computation on the prompt), focusing 100% of backpropagation on generating the target modality.

#3. Experimental Results: Five Generation Modes from 12.51M Parameters

Keeping our compact 6-layer, 6-head, 384-hidden-dimension configuration (while expanding the context window block_size to 512 to fit complete music-text pairs), the model totals 12.51M parameters.

After joint training for 2,000 steps on a mix of Chinese Journey to the West (650k tokens), an English novel (500k tokens), 12 Chopin piano pieces (63k MIDI tokens), and 488 paired "music clip — bilingual description" samples, a single set of weights successfully handles five tasks:

  1. Chinese Continuation: Given [TOK_TEXT] + "悟空道：" ("Wukong said:"), the model automatically stays within the Chinese character range and generates dialogue in the style of Journey to the West;
  2. English Continuation: Given [TOK_TEXT] + "Jack looked", the model automatically switches to English vocabulary and grammar distributions;
  3. Unconditional Chopin Performance: Given only a single start token [TOK_MUSIC], the model autoregressively generates the following Chopin-style piano passage:
  1. Text-to-Music ($\text{Text} \to \text{Music}$): Given a Chinese mood and tempo description ending with [TOK_MUSIC], Self-Attention directly attends across modalities to the preceding Chinese text features and generates a matching piano MIDI sequence;
  2. Music-to-Text ($\text{Music} \to \text{Text}$): Given an unseen validation piano MIDI sequence ending with [TOK_TEXT], the model attends back over the pitch contours and temporal density of the entire note sequence to output a matching Chinese musical commentary.

#4. Summary: "Multimodal Understanding" Is Still Conditional Probability in a Shared Space

Connecting the experiments across these first three posts reveals how simple the underlying logic of native multimodal LLMs really is:

  1. Modality Switching Is Subspace Routing: There is no if modality == "music" branching logic anywhere inside the model. Every Chinese character, English letter, and piano key shares the exact same 6 layers of Attention and MLP weights. The reason it never randomly emits piano notes while writing English is purely that the training data drives the conditional probability of MIDI tokens (0..307) given a TOK_TEXT context close to zero (and applying a simple token_mask at inference time can strictly lock the output modality).
  2. Cross-Modal Alignment Is Cross-Range Attention Retrieval: During Text-to-Music or Music-to-Text generation, QKV Self-Attention performs the exact same operation as in monolingual translation—except the Query comes from the text token range (312..4887) while the Keys and Values being retrieved live in the MIDI token range (0..307).
  3. From "Working Prototype" to Modern Architectures: At this point, we have pushed the 2019 GPT-2 architecture to its limit. Yet as foreshadowed at the end of Part 1, when scaling to longer contexts and larger parameter counts, standard GPT-2 remains bottlenecked by learned absolute positional embeddings (wpe), dense MLP compute coupling, and inference without a KV Cache—the exact modern LLM architectural evolutions we will unpack starting in Part 4 (Tokenizers and Positional Encoding).

See Also: