Skip to content

feat(vlm): add Kimi-K3 pokemon ShareGPT SFT recipe with collate and dataset - #1

Open
xiaoliang0601 wants to merge 3 commits into
kimi-k3from
huiyingl/feat/k3-vlm-pokemon-sft
Open

xiaoliang0601 wants to merge 3 commits into
kimi-k3from
huiyingl/feat/k3-vlm-pokemon-sft

Conversation

@xiaoliang0601

Copy link
Copy Markdown
Owner

No description provided.

xiaoliang0601 and others added 2 commits September 15, 2026 17:21
…ded save

Parameters that never receive a gradient (e.g. Kimi-K3 layer-0
self_attention_res_* params, whose first residual mix is skipped while the
block residual is still empty) never get lazily-created Adam optimizer state.
get_optimizer_state_dict() then silently omits their state at save time, yet
the optimizer expects it at resume, so DCP load fails with a missing
"...weight.step" key. Mirror the native (non-sharded) path and materialize
zero state for all params before building the sharded optimizer state dict.

Co-Authored-By: Claude <noreply@anthropic.com>
Signed-off-by: root <xuezunyao@xiaohongshu.com>
…ataset

Enable Kimi-K3 VLM SFT on the pokemon-gpt4o Chinese caption set:

- kimi_k3_vl_collate_fn: feeds the model a single <|media_pad|> per image
  (K3's forward expands it internally) and returns merged-length, pre-shifted
  labels for the recipe's no-shift MaskedCrossEntropy. Right-pads to
  batch_max+1 so left_padding is always False and the internal image-token
  merge layout is deterministic.
- make_pokemon_sharegpt_dataset: loads the local ShareGPT parquet, decodes
  inline image bytes lazily, and maps human/gpt turns to the user/assistant
  conversation structure the collate expects.
- k3_pokemon_zh.yaml: FSDP2 + EP8 recipe (K3 has no TP), 200-step run against
  the real 5L/128E MXFP4 checkpoint.

Co-Authored-By: Claude <noreply@anthropic.com>
Signed-off-by: root <xuezunyao@xiaohongshu.com>
@xiaoliang0601
xiaoliang0601 force-pushed the huiyingl/feat/k3-vlm-pokemon-sft branch from 0751bd9 to 4fef284 Compare September 15, 2026 18:09
torch.optim.AdamW keeps no internal fp32 master, so the resident parameter
is the master copy. bf16 storage rounds Adam's ~lr-sized updates below the
bf16 ULP (swamping), stalling early descent; float32 storage captures them
and matches Relax's precision-aware fp32-master optimizer. Also drop warmup
to 0 to align with the Relax reference (Megatron default). FSDP2 still
computes/all-gathers in bf16 via mp_policy.

Co-Authored-By: Claude <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant