vanbasten23’s gists

vanbasten23 / gist:abc3b2db22748e1f16aa0e92c6b91404

Created December 4, 2025 18:29

	exec ${PAGER:-/usr/bin/less -R} "$0" \|\| exit 1

	Test settings: forge with network access
	Host details: itmm4.prod.google.com Linux 6.6.65-smp-1300.170.0.0 x86_64 astoria-genoa-base
	executor.INFO: analog/view?storage=borgremote&bns=/bns/it/borg/it/bns/build-forge-executor-tpu/prod-cbf-ghostlite.forge-executor/0&min_time=1764872604000000&ts=1764872614000000
	Test command:
	cd /build/work/aef67bf50706fee86777a93cc065340a246c/google3/runfiles/google3 && \
	env - \
	BORG_CELL=it \
	CUSTOM_METRICS_DIR=/build/work/aef67bf50706fee86777a93cc065340a246c/google3/../custom_metrics \

vanbasten23 / gist:42e056748dabb38e81d34bb3d3cc19b3

Created November 19, 2025 18:31

	Let's trace the values for my_id = 1 with num_devices = 4:

	outer_step phase Accumulation Source left_copy_device right_copy_device Device providing the data
	0 LEFT x_ref[left_copy_device, ...] (1+0+1)%4 = 2 (1-0-1)%4 = 0 Device 2
	0 RIGHT x_ref[right_copy_device, ...] (1+0+1)%4 = 2 (1-0-1)%4 = 0 Device 0
	1 LEFT x_ref[left_copy_device, ...] (1+1+1)%4 = 3 (1-1-1)%4 = 3 Device 3
	1 RIGHT x_ref[right_copy_device, ...] (1+1+1)%4 = 3 (1-1-1)%4 = 3 Device 3
	2 LEFT x_ref[left_copy_device, ...] (1+2+1)%4 = 0 (1-2-1)%4 = 2 Device 0
	2 RIGHT x_ref[right_copy_device, ...] (1+2+1)%4 = 0 (1-2-1)%4 = 2 Device 2
	As you can see, with each outer_step, the *_copy_device variables change, ensuring that the reduction operation fetches data from a new, distinct device. This systematic progression guarantees that by the end of all steps, each device has accumulated its required portion of the total sum from all other devices.

vanbasten23 / gist:23096fb5a117e09587a5d5f34232ae7d

Created November 14, 2025 19:01

	import jax
	from jax import export
	import jax.numpy as jnp
	import pickle
	import time
	import statistics

	with open("/home/xiowei_google_com/new_exports.pkl", "rb") as f:
	data = pickle.load(f)

vanbasten23 / gist:1dd9a30a26fda4d7d112cf7611bcbaa4

Created November 14, 2025 19:00

	import jax
	from jax import export
	import jax.numpy as jnp
	import pickle
	import time
	import statistics

	with open("/home/xiowei_google_com/old_exports.pkl", "rb") as f:
	data = pickle.load(f)

vanbasten23 / gist:4a983ec0b95ce28a823b8217193ad1f8

Created November 14, 2025 18:29

	1. Start the benchmark server in vscode as [this](https://gist.github.com/vanbasten23/dd4f3cbb314a7b9cf6c003103c23c019). Select the correct python intepreter.
	2. Then start the vllm server in debugger.
	3. After the server is up and running.
	4. Add the breakpoint (remember to turn of dynamo and jax jit)
	5. Use the [script](https://gist.github.com/vanbasten23/726b28f072993fb7587482672b9c96a9) to send benchmarking request. Make sure to use the correct conda/python.
	6. Then dump the input and output.


	=========================
	pip install flatbuffers

vanbasten23 / gist:726b28f072993fb7587482672b9c96a9

Created November 14, 2025 18:18

	#!/bin/bash

	# Usage:
	# bash run_tpu_benchmark_client.sh --model Qwen/Qwen2.5-1.5B-Instruct --tp 1

	LONGOPTS=model:,tp:,profile
	# Parse arguments
	PARSED=$(getopt --options=$OPTIONS --longoptions=$LONGOPTS --name "$0" -- "$@")
	if [[ $? -ne 0 ]]; then
	exit 2

vanbasten23 / gist:dd4f3cbb314a7b9cf6c003103c23c019

Created November 14, 2025 18:16

	{
	"name": "newjax_benchmark_server",
	"type": "debugpy",
	"request": "launch",
	"program": "/home/xiowei_google_com/miniconda3/envs/vllm_newjax/bin/vllm",
	"console": "integratedTerminal",
	"justMyCode": false,
	"env": {
	"MODEL_IMPL_TYPE": "vllm",
	"TPU_BACKEND_TYPE": "jax",

vanbasten23 / gist:7618a394435a024594fd58de57fa337d

Created November 14, 2025 06:24

	import jax
	from jax import export
	import jax.numpy as jnp
	import pickle
	import time
	import statistics

	with open("/home/xiowei_google_com/old_exports.pkl", "rb") as f:
	data = pickle.load(f)

vanbasten23 / gist:f3b509b4f1c05d9e4a93e0c42de97a2c

Created October 31, 2025 20:20

	local keymap = vim.keymap.set
	local opts = { noremap = true, silent = true }

	-- remap leader key
	keymap("n", "<Space>", "", opts)
	vim.g.mapleader = " "
	vim.g.maplocalleader = " "

	-- yank to system clipboard
	keymap({"n", "v"}, "<leader>y", '"+y', opts)

vanbasten23 / gist:eec91565e61808da8e726067732bb127

Created October 31, 2025 18:07

	from transformers import AutoTokenizer, AutoModelForCausalLM
	from peft import PeftModel
	import torch


	base = "Qwen/Qwen2.5-3B-Instruct"
	adapter = "./lora-1plus1-666"
	tok = AutoTokenizer.from_pretrained(base)
	m = AutoModelForCausalLM.from_pretrained(base, torch_dtype=torch.bfloat16).to("cuda" if torch.cuda.is_available() else "cpu")
	m = PeftModel.from_pretrained(m, adapter)

XiongfeiWei vanbasten23