Skip to content

Snehalv dsv4 conv utility#4078

Draft
snehalv2002 wants to merge 15 commits into
mainfrom
snehalv-dsv4-conv-utility
Draft

Snehalv dsv4 conv utility#4078
snehalv2002 wants to merge 15 commits into
mainfrom
snehalv-dsv4-conv-utility

Conversation

@snehalv2002

@snehalv2002 snehalv2002 commented Jun 5, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Description

Standalone and conversion util support for DeepSeek V4 Scanned.
FIXES: b/509930555

Tests

Checklist

Before submitting this PR, please make sure (put X in square brackets):

  • I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
  • I have necessary comments in my code, particularly in hard-to-understand areas.
  • I have run end-to-end tests tests and provided workload links above if applicable.
  • I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

@snehalv2002 snehalv2002 force-pushed the snehalv-dsv4-conv-utility branch from a201455 to e45a786 Compare June 15, 2026 19:50
@codecov

codecov Bot commented Jun 15, 2026 •
edited
Loading

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 10.16043% with 168 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
...xtext/checkpoint_conversion/utils/param_mapping.py 1.81% 108 Missing ⚠️
...rc/maxtext/checkpoint_conversion/utils/hf_shape.py 0.00% 38 Missing ⚠️
src/maxtext/checkpoint_conversion/to_maxtext.py 0.00% 10 Missing ⚠️
src/maxtext/checkpoint_conversion/utils/utils.py 0.00% 9 Missing ⚠️
src/maxtext/models/deepseek.py 33.33% 1 Missing and 1 partial ⚠️
src/maxtext/layers/attention_mla.py 50.00% 0 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

@snehalv2002 snehalv2002 force-pushed the snehalv-dsv4-conv-utility branch 2 times, most recently from 278f2ec to b210f93 Compare June 17, 2026 16:57
snehalv2002 added 14 commits June 23, 2026 07:02
@snehalv2002
feat: DeepSeek V4 checkpoint conversion and validation utilities
c8ff8a0 
- Add standalone script for scanned DeepSeek V4 checkpoint conversion.
- Add DeepSeek V4 unscanned checkpoint conversion script and validation tool.
- Refactor checkpoint conversion to cast path keys to strings and add test script for unscanned paths.

TAG=agy
CONV=6d130342-8f6e-4abd-a4c5-dbbfd7180ae5
@snehalv2002
update conversion script to fit new dsv4 implementation
79b4936
@snehalv2002
feat: add DeepSeek-V4 parameter inspection scripts and config updates…
350a4e4 
… for checkpoint conversion.
@snehalv2002
feat: add support for DeepSeek-V4 checkpoint conversion and inspectio…
8b584c9 
…n tools
@snehalv2002
Fix DeepSeek-V4 scanned and unscanned checkpoint conversion mappings …
283e158 
…and shapes

- Updated SCANNED_LAYER_MAPPING to include standalone layers (0, 1, 2) and alternate scanned block ranges.
- Corrected query projection (wq_b) shapes for standard and indexer parameters.
- Updated hf_shape registry with dynamic compressor key-value projection and normalization shapes.
- Successfully verified both scanned and unscanned conversions via dry-run shape matching.

TAG=agy
CONV=6d130342-8f6e-4abd-a4c5-dbbfd7180ae5
@snehalv2002
Update checkpoint_xpk.sh topology parameters and scanned checkpoint path
9075733 
- Updated cluster name to mlperf-v5p.
- Updated zone to europe-west4-b.
- Updated tpu-type to v5p-256.
- Updated load_parameters_path to the verified scanned DeepSeek-V4 checkpoint.
- Updated model_name to deepseek4-284b.

TAG=agy
CONV=6d130342-8f6e-4abd-a4c5-dbbfd7180ae5
@snehalv2002
Configure checkpoint_xpk.sh with manifest patching pipeline for reser…
c2f05d7 
…ved cluster taints

- Add dry-run manifest generation.
- Add fix_manifest.py patch step to configure reservation and toleration.
- Add kubectl apply step to submit workload.

TAG=agy
CONV=6d130342-8f6e-4abd-a4c5-dbbfd7180ae5
@snehalv2002
Update checkpoint_xpk.sh to use real image upload and manifest patchi…
3706c06 
…ng flow

TAG=agy
CONV=6d130342-8f6e-4abd-a4c5-dbbfd7180ae5
@snehalv2002
Optimize checkpoint_xpk.sh pipeline order to avoid workload creation …
6a99409 
…conflicts

TAG=agy
CONV=6d130342-8f6e-4abd-a4c5-dbbfd7180ae5
@snehalv2002
Update checkpoint_xpk.sh to run train_compile.py with compile configu…
4132106 
…rations

TAG=agy
CONV=6d130342-8f6e-4abd-a4c5-dbbfd7180ae5
@snehalv2002
Update checkpoint_xpk.sh to run train.py training script
fbffe53 
TAG=agy
CONV=6d130342-8f6e-4abd-a4c5-dbbfd7180ae5
@snehalv2002
Ignore scratch-pad/ folder
c0870d2 
TAG=agy
CONV=6d130342-8f6e-4abd-a4c5-dbbfd7180ae5
@snehalv2002
Align finalize function in unscanned conversion script with scanned s…
7409338 
…cript

TAG=agy
CONV=6d130342-8f6e-4abd-a4c5-dbbfd7180ae5
@snehalv2002
Fix assertions and config for local CPU test
eeda017
@snehalv2002 snehalv2002 force-pushed the snehalv-dsv4-conv-utility branch from 5bb1ba6 to eeda017 Compare June 23, 2026 07:23
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@parambole parambole Awaiting requested review from parambole parambole will be requested when the pull request is marked ready for review parambole is a code owner
@shuningjin shuningjin Awaiting requested review from shuningjin shuningjin will be requested when the pull request is marked ready for review shuningjin is a code owner
@RissyRan RissyRan Awaiting requested review from RissyRan RissyRan will be requested when the pull request is marked ready for review RissyRan is a code owner
@suexu1025 suexu1025 Awaiting requested review from suexu1025 suexu1025 will be requested when the pull request is marked ready for review suexu1025 is a code owner
@jiangjy1982 jiangjy1982 Awaiting requested review from jiangjy1982 jiangjy1982 will be requested when the pull request is marked ready for review jiangjy1982 is a code owner
@gobbleturk gobbleturk Awaiting requested review from gobbleturk gobbleturk will be requested when the pull request is marked ready for review gobbleturk is a code owner
@bvandermoon bvandermoon Awaiting requested review from bvandermoon bvandermoon will be requested when the pull request is marked ready for review bvandermoon is a code owner
@gagika gagika Awaiting requested review from gagika gagika will be requested when the pull request is marked ready for review gagika is a code owner
@shralex shralex Awaiting requested review from shralex shralex will be requested when the pull request is marked ready for review shralex is a code owner
@richjames0 richjames0 Awaiting requested review from richjames0 richjames0 will be requested when the pull request is marked ready for review richjames0 is a code owner
@NuojCheng NuojCheng Awaiting requested review from NuojCheng NuojCheng will be requested when the pull request is marked ready for review NuojCheng is a code owner
@xibinliu xibinliu Awaiting requested review from xibinliu xibinliu will be requested when the pull request is marked ready for review xibinliu is a code owner
@huytransformer huytransformer Awaiting requested review from huytransformer huytransformer will be requested when the pull request is marked ready for review huytransformer is a code owner
@aireenmei aireenmei Awaiting requested review from aireenmei aireenmei will be requested when the pull request is marked ready for review aireenmei is a code owner
@igorts-git igorts-git Awaiting requested review from igorts-git igorts-git will be requested when the pull request is marked ready for review igorts-git is a code owner
@hengtaoguo hengtaoguo Awaiting requested review from hengtaoguo hengtaoguo will be requested when the pull request is marked ready for review hengtaoguo is a code owner
@abhinavclemson abhinavclemson Awaiting requested review from abhinavclemson abhinavclemson will be requested when the pull request is marked ready for review abhinavclemson is a code owner
@khatwanimohit khatwanimohit Awaiting requested review from khatwanimohit khatwanimohit will be requested when the pull request is marked ready for review khatwanimohit is a code owner
@vipannalla vipannalla Awaiting requested review from vipannalla vipannalla will be requested when the pull request is marked ready for review vipannalla is a code owner
@SurbhiJainUSC SurbhiJainUSC Awaiting requested review from SurbhiJainUSC SurbhiJainUSC will be requested when the pull request is marked ready for review SurbhiJainUSC is a code owner
@A9isha A9isha Awaiting requested review from A9isha A9isha will be requested when the pull request is marked ready for review A9isha is a code owner
@dipannita08 dipannita08 Awaiting requested review from dipannita08 dipannita08 will be requested when the pull request is marked ready for review dipannita08 is a code owner
@darisoy darisoy Awaiting requested review from darisoy darisoy will be requested when the pull request is marked ready for review darisoy is a code owner

At least 2 approving reviews are required to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@snehalv2002