Cameron R. Wolfe wolfecameron

Research @ Netflix • Writer of Deep (Learning) Focus • PhD in CS from Rice University

wolfecameron / expert_layer.py

Created March 6, 2025 21:17

PyTorch implementation of a feed-forward expert layer within an MoE.

	"""
	Based upon ColossalAI OpenMoE
	"""

	from torch import nn

	class MOELayer(nn.Module):
	def __init__(
	self,
	d,

wolfecameron / moe_block.py

Created March 6, 2025 22:24

MoE block for an MoE-based decoder-only transformer model in PyTorch.

wolfecameron / bidir_self_attn.py

Last active March 29, 2025 17:36

An implementation of bidirectional self-attention in PyTorch.

	import math
	import torch
	from torch import nn
	import torch.nn.functional as F

	class SelfAttention(nn.Module):

	def __init__(self, d):
	"""
	Arguments:

wolfecameron / cross_attention.py

Last active November 6, 2025 02:38

An implementation of cross-attention in PyTorch.

	import math
	import torch
	from torch import nn
	import torch.nn.functional as F

	class CrossAttention(nn.Module):

	def __init__(self, d):
	"""
	Arguments:

wolfecameron / olmo_trace_index.py

Created May 17, 2025 04:30

Build an infini-gram index (and store other relevant info) for use in OLMoTrace.

	import os
	import json
	from collections import Counter
	import tempfile

	from transformers import AutoTokenizer

	# load tokenizer / data
	enc = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf", add_bos_token=False, add_eos_token=False)
	data_rows = [{'text': 'here is some training data'}, ...]

wolfecameron / olmo_trace.py

Last active June 19, 2025 23:35

Tracing text using the algorithm proposed by OLMoTrace: https://arxiv.org/abs/2504.07096

	import ast
	import math
	import random

	from infini_gram.engine import InfiniGramEngine
	from transformers import AutoTokenizer

	def compute_longest_prefix(query, doc):
	"""helper function for computing longest prefix of query that exists
	within a document"""