wire Fireworks distillation smoke test path

[yi-fireworks]

Summary

• wire on-policy distillation through a new Fireworks-specific runtime so a CPU-only orchestrator can create the trainer client, attach or create the rollout deployment, hotload weights, and run the recipe against remote Fireworks infrastructure
• adapt the distillation path to the public Fireworks training SDK surface, including tokenizer resolution for Fireworks model ids and the currently available WeightSyncer API
• leave focused instrumentation in the first-batch path so the remaining stall is pinned to the trainer forward_backward_async(...) call

Current blocker / help wanted

The smoke test now gets through rollout sampling, minibatch assembly, teacher KL, and deployment hotload, but the actual training RPC still hangs.

What we observed:

• rollout sampling succeeds
• minibatch assembly succeeds
• teacher KL succeeds
• initial hotload and warmup succeed
• the run hangs on await training_client.forward_backward_async(...)

What we ruled out:

• same hang with the original pipelined train step and with a serial Fireworks-specific train step
• same hang with loss_fn=cross_entropy and kl_penalty_coef=0
• teacher KL is not the blocker

While the run is stuck, the trainer job still reports:

• JOB_STATE_RUNNING
• status OK
• progress counters remain zero

If anyone familiar with this trainer path can help, the main question is whether this endpoint / SDK path is expected to support forward_backward_async in this configuration, or whether there is a known issue with this shape / model / lora_rank=0 setup.

Repro command

Full distillation-path repro:

source /home/yi/venv/bin/activate && export PYTHONUNBUFFERED=1 && export FIREWORKS_ACCOUNT_ID=pyroworks && python -m tinker_cookbook.recipes.distillation.on_policy_distillation \
  base_url="https://trainer-pyroworks-rlor-kltlnlct82fw9z1d.us-ohio-1.direct.fireworks.ai:30443/" \
  fireworks_base_model_name="accounts/fireworks/models/qwen3-4b-instruct-2507" \
  model_name="accounts/fireworks/models/qwen3-4b-instruct-2507" \
  teacher_model="accounts/fireworks/models/qwen3-4b-instruct-2507" \
  renderer_name=qwen3_instruct \
  lora_rank=0 \
  dataset=deepmath \
  max_steps=1 \
  groups_per_batch=1 \
  group_size=1 \
  max_tokens=64 \
  eval_every=0 \
  save_every=1

Control repro that still hangs even without distillation loss / KL:

source /home/yi/venv/bin/activate && export PYTHONUNBUFFERED=1 && export FIREWORKS_ACCOUNT_ID=pyroworks && python -m tinker_cookbook.recipes.distillation.on_policy_distillation \
  base_url="https://trainer-pyroworks-rlor-kltlnlct82fw9z1d.us-ohio-1.direct.fireworks.ai:30443/" \
  fireworks_base_model_name="accounts/fireworks/models/qwen3-4b-instruct-2507" \
  model_name="accounts/fireworks/models/qwen3-4b-instruct-2507" \
  teacher_model="accounts/fireworks/models/qwen3-4b-instruct-2507" \
  renderer_name=qwen3_instruct \
  lora_rank=0 \
  dataset=deepmath \
  max_steps=1 \
  groups_per_batch=1 \
  group_size=1 \
  max_tokens=64 \
  eval_every=0 \
  save_every=1 \
  loss_fn=cross_entropy \
  kl_penalty_coef=0

Test plan

• verify imports and py_compile on touched files
• create trainer / deployment and run real smoke-test attempts against remote Fireworks infra
• get a successful 1-step forward_backward_async through the remote trainer

Made with Cursor