🔬Modal Research Exchange

{
  "base_model": "Qwen/Qwen3-4B",
  "baseline": "Same fork, data, and 100 sparse LoRA steps with indexer KL during training, but no 50-step indexer warmup. The indexer starts from random init. Dense-attention LoRA is logged as an upper-bound reference.",
  "batch_size": 64,
  "budget": {
    "gpu_hours": 2.0,
    "max_wall_minutes": 120
  },
  "command": "python - <<'EOF'\n# custom: indexer warmup + sparse attention come from the forked Spindle; loop is stock tinker SDK\nimport os, tinker\nsc = tinker.ServiceClient(base_url=os.environ['TINKER_BASE_URL'])\ntc = sc.create_lora_training_client(base_model='Qwen/Qwen3-4B', rank=32)\nbatches = load_pg19_chunks(seq_len=4096, batch_size=64)  # user helper, fixed seed\nWARMUP = int(os.environ.get('WARMUP', 50))  # 0 for baseline arm\nfor step, b in enumerate(batches[:WARMUP + 100]):\n    loss = 'indexer_kl' if step < WARMUP else 'cross_entropy'  # 'indexer_kl' added by fork\n    tc.forward_backward(b, loss).result()\n    tc.optim_step(tinker.AdamParams(learning_rate=1e-4)).result()\ntc.save_weights_for_sampler(name=f'sparse-idx-warmup{WARMUP}').result()\nEOF",
  "dataset": {
    "hf_path": "deepmind/pg19",
    "n_examples": 9600,
    "name": "PG19",
    "split": "train"
  },
  "eval": {
    "benchmark": "PG19 validation, 4096-token chunks, evaluated with sparse top-512 attention",
    "metric": "mean next-token cross-entropy (nats), plus top-k recall of dense attention mass",
    "n_examples": 200
  },
  "hypothesis": "Warming up a top-k indexer by KL distillation from frozen full attention, then LoRA-adapting under sparse attention, gives held-out PG19 loss at least 0.05 nats lower than the same sparse LoRA run with a randomly initialized indexer.",
  "lora_rank": 32,
  "max_seq_len": 4096,
  "method": "Needs a Spindle fork that adds a DSA-style lightning indexer to Qwen/Qwen3-4B. Qwen3-4B is used because it is pure full attention, while Qwen3.5 uses hybrid linear attention. The indexer is a few low-dim heads per layer that score past tokens. The fork exposes three things: (a) an indexer-KL loss against the aggregated full-attention distribution, with the main model frozen; (b) top-k (k=512) token-selection sparse attention in forward and backward; (c) the indexer as a trainable parameter alongside the LoRA weights. Treatment stage 1 runs 50 steps that train only the indexer with KL to dense attention, using dense forward and a frozen base. Treatment stage 2 runs 100 steps of LoRA rank 32 next-token SFT on 4096-token PG19 chunks under top-k sparse attention. During stage 2 the indexer keeps its KL loss on selected tokens, detached from the LM loss, as in DSA. Both arms use the same data order, seed and LR. A dense-attention LoRA run with the same 100 steps is logged as a reference ceiling, not as an arm. Neither linked source could be read, so the design follows the Slack description and the DeepSeek/GLM indexer recipe.",
  "optimizer": {
    "lr": 0.0001,
    "name": "adam"
  },
  "steps": 150,
  "success_criterion": "Validated if treatment eval loss is at least 0.05 nats below baseline and within 0.05 nats of the dense LoRA reference. Seed noise is to be checked to be under 0.02 nats via a second baseline seed. Failed if the gap to baseline is under 0.02 nats.",
  "title": "Indexer warmup from full attention versus cold-start indexer for sparse LoRA adaptation",
  "training_type": "sft"
}

batch_size

2

max_seq_len

4,096

lora_rank

32

topk

512

eval_chunks

200

eval_tokens_per_arm

819,200

local_tests_passed

185

local_tests_skipped

1

treatment_training_steps

150

total_training_steps

450

treatment_eval_loss_nats

3.478

treatment_topk_dense_mass_recall

0.8108

baseline_seed0_eval_loss_nats

3.608

baseline_seed0_topk_dense_mass_recall

0.776

baseline_seed1_eval_loss_nats

3.628

baseline_seed1_topk_dense_mass_recall

0.7785

dense_reference_eval_loss_nats

3.063

baseline_minus_treatment_nats

0.1307

treatment_minus_dense_nats

0.4146

baseline_seed_difference_nats

0.01962

completed_arms

4

elapsed_seconds

3,862.2

remaining_run_apps_active

0

remaining_run_containers

0

remaining_run_volumes

0

remaining_run_named_dicts

0

observed_proxy_401_count

0