# Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. import torch import torch.nn as nn from torch.nn import CrossEntropyLoss from collections import namedtuple from transformers.models.gpt2 import GPT2LMHeadModel Outputs = namedtuple("Outputs", ["loss", "inputs_embeds", "logits"]) MAX_N_LATENT = 8 class Coconut(nn.Module): def __init__( self, base_causallm, latent_token_id, start_latent_id, end_latent_id, eos_token_id, backprop_depth=None, ): super(Coconut, self).__init__() self.gen_forward_cnt = 0 self.base_causallm = base_causallm self.latent_token_id = latent_token_id self.eos_token_id = eos_token_id self.start_latent_id = start_latent_id self.end_latent_id = end_latent_id # Truncated BPTT window: if set to W, gradients only flow through the last W # latent recurrence steps (+ the final answer pass); older latent states and # their KV cache are detached. This keeps backward compute + activation memory # O(W) instead of O(depth), enabling higher graph depth at ~constant latency. # None => full BPTT (upstream behavior). self.backprop_depth = backprop_depth # tested with GPT2 and Llama3 if isinstance(self.base_causallm, GPT2LMHeadModel): self.embedding = self.base_causallm.transformer.get_input_embeddings() else: self.embedding = self.base_causallm.get_input_embeddings() def forward(self, input_ids, attention_mask, labels, position_ids, **kwargs): logits = [] latent_indices = ( input_ids == self.latent_token_id ).nonzero() # (num_latent_tokens_in_the_batch, 2) # Group latent positions per batch instance. The upstream version did this with # a nested Python loop that called `.item()` and compared GPU scalars for every # (instance x latent) pair -> O(bs^2 * depth) GPU->CPU syncs per forward, which # (not the model compute or backprop) was the dominant cost and made deep latent # chains explode. We transfer the index pairs to host ONCE and group in pure Python. latent_lists = [[] for _ in range(input_ids.shape[0])] for r, c in latent_indices.tolist(): latent_lists[r].append(c) # nonzero() already returns row-major order max_n_latents = max([len(l) for l in latent_lists]) next_compute_range = (0, input_ids.shape[1]) inputs_embeds = self.embedding(input_ids) if max_n_latents > 0: next_compute_range = (0, latent_indices[:, 1].min().item()) # before the earliest latent token position kv_cache = None for pass_idx in range(max_n_latents): if kv_cache == None: # first forward pass outputs = self.base_causallm( inputs_embeds=inputs_embeds[ :, next_compute_range[0] : next_compute_range[1], : ], attention_mask=attention_mask[ :, next_compute_range[0] : next_compute_range[1] ], position_ids=position_ids[ :, next_compute_range[0] : next_compute_range[1] ], output_hidden_states=True, ) hidden_states_offset = 0 else: # extract kv cache to reuse past_key_values = [ ( k[:, :, : next_compute_range[0], :], v[:, :, : next_compute_range[0], :], ) for k, v in kv_cache ] outputs = self.base_causallm( inputs_embeds=inputs_embeds[ :, next_compute_range[0] : next_compute_range[1], : ], attention_mask=attention_mask[:, : next_compute_range[1]], position_ids=position_ids[ :, next_compute_range[0] : next_compute_range[1] ], past_key_values=past_key_values, output_hidden_states=True, ) hidden_states_offset = next_compute_range[0] # when we use kv_cache for the first k tokens # in `outputs.hidden_states`, [0, k) will be skipped # so we need to keep this offset to correctly use the last hidden states logits.append(outputs.logits) next_compute_range = ( next_compute_range[1], ( input_ids.shape[1] if pass_idx + 1 >= max_n_latents else next_compute_range[1] + 1 ), ) hidden_states = outputs.hidden_states[ -1 ] # Get the last layer hidden states kv_cache = outputs.past_key_values # Truncated BPTT: for passes older than `backprop_depth` steps from the end, # detach the fed-back thought and the KV cache so the autograd graph only # retains the last W recurrent steps. `passes_remaining` counts steps after # this one; when it is >= W this pass falls outside the window. if self.backprop_depth is not None: passes_remaining = max_n_latents - 1 - pass_idx if passes_remaining >= self.backprop_depth: hidden_states = hidden_states.detach() if kv_cache is not None: kv_cache = [ (k.detach(), v.detach()) for (k, v) in kv_cache ] # feedback the continuous thoughts to the input_embeds # first decide the positions to feedback filling_indices = [ (instance_idx, mask_list[pass_idx]) for instance_idx, mask_list in enumerate(latent_lists) if len(mask_list) > pass_idx ] # Vectorized continuous-thought feedback (equivalent to the upstream # per-element rebuild, but ~batch*seq times fewer GPU ops). The upstream # code materialized batch*seq individual 1-D tensor views in Python every # latent pass, which dominates runtime on this small synthetic model and # scales with the number of latents. We instead clone once (to stay # autograd-safe / out-of-place w.r.t. the tensor still needed for backward) # and scatter the preceding hidden states into the latent positions. if filling_indices: b_idx = torch.tensor( [b for b, _ in filling_indices], device=inputs_embeds.device ) t_idx = torch.tensor( [t for _, t in filling_indices], device=inputs_embeds.device ) src = hidden_states[b_idx, t_idx - 1 - hidden_states_offset, :] inputs_embeds = inputs_embeds.clone() inputs_embeds[b_idx, t_idx, :] = src # final pass outputs = self.base_causallm( inputs_embeds=inputs_embeds[ :, next_compute_range[0] : next_compute_range[1], : ], attention_mask=attention_mask[:, : next_compute_range[1]], position_ids=position_ids[:, next_compute_range[0] : next_compute_range[1]], past_key_values=( [ ( k[:, :, : next_compute_range[0], :], v[:, :, : next_compute_range[0], :], ) for k, v in kv_cache ] if kv_cache else None ), output_hidden_states=True, ) logits.append(outputs.logits) self.gen_forward_cnt += max_n_latents + 1 logits = torch.cat(logits, dim=-2) shift_logits = logits[..., :-1, :].contiguous() shift_labels = labels[..., 1:].contiguous() loss_fct = CrossEntropyLoss() loss = loss_fct( shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1) ) return Outputs(loss=loss, inputs_embeds=inputs_embeds, logits=logits) def train(self): self.base_causallm.train() def eval(self): self.base_causallm.eval() def generate( self, input_ids, attention_mask, # attention_mask is not used max_new_tokens=16, output_embedding=False, synced_gpus=False, **kwargs ): self.gen_forward_cnt = 0 assert input_ids.shape[0] == 1, "only support batch_size == 1 now" tokens = input_ids[0].detach().tolist() labels = input_ids.clone() # placeholder. not used. outputs = self.forward( input_ids, torch.ones_like(input_ids, device=input_ids.device), labels, torch.arange( 0, input_ids.shape[1], dtype=torch.long, device=input_ids.device ).reshape(1, -1), ) inputs_embeds = outputs.inputs_embeds # get the first token using the current hidden state next_token = torch.argmax(outputs.logits[0, -1]).item() tokens.append(next_token) new_token_embed = self.embedding( torch.tensor(next_token, device=input_ids.device) ).view(1, 1, -1) new_inputs_embeds = torch.cat((inputs_embeds, new_token_embed), dim=1) # get other tokens for _ in range(max_new_tokens - 1): outputs = self.base_causallm(inputs_embeds=new_inputs_embeds) self.gen_forward_cnt += 1 next_token = torch.argmax(outputs.logits[0, -1]).item() if next_token == self.eos_token_id: break tokens.append(next_token) new_token_embed = self.embedding( torch.tensor(next_token, device=input_ids.device) ).view(1, 1, -1) new_inputs_embeds = torch.cat((new_inputs_embeds, new_token_embed), dim=1) if synced_gpus: # in FSDP, the number of forward pass need to be the same across devices while ( self.gen_forward_cnt < max_new_tokens + MAX_N_LATENT ): # leave some room for latent tokens self.gen_forward_cnt += 1 _ = self.base_causallm(inputs_embeds=new_inputs_embeds) if output_embedding: # for analysis purpose return torch.tensor(tokens).view(1, -1), new_inputs_embeds else: return torch.tensor(tokens).view(1, -1)