added example for bidirectional checkpoint testing #1540

wesleytruong · 2025-08-06T18:57:24Z

This pr adds

an example script for bidirectional testing of checkpoint conversion scripts
a checkpoint_conversion.md to describe our methodology.

tianyu-l · 2025-08-07T01:14:34Z

scripts/checkpoint_conversion/example.py

+    logger.info(f"Loading chkpt at: {checkpoint_path}")
+    load_from_hf = False
+    for filename in os.listdir(checkpoint_path):
+        if filename == "model.safetensors.index.json":


this is not reliable -- if there is only one .safetensors file, there won't be such index file

tianyu-l · 2025-08-07T01:15:33Z

scripts/checkpoint_conversion/checkpoint_conversion.md

rename it to README.md (so it's displayed when entering this folder)

tianyu-l · 2025-08-07T01:17:21Z

scripts/checkpoint_conversion/checkpoint_conversion.md

+### Sanity Check (Greedy Decode)
+A quick way to sanity check if your conversion is correct is to perform greedy decoding inference on both the initial and converted checkpoints and confirm that they are the same. This method doesn't guarantee correctness but will very likely result in a fast **true negative** if the model definitions are not the same. For greedy decoding, the `generation/test_generate.py` script can be used.
+
+Note that the model definitions can be influenced by external factors than correctness of weight conversion. For example, using our verified `convert_to_hf.py` script then running greedy decoding using HF `transformers` without a correct `config.json` will result in a **false negative** since our weights are correct but the model definition is incorrect due to `config.json`.


This is too obscure than it needs to be. Don't need say what could go wrong, say what they need to do to get right.
E.g. you can just say in order to use HF transformers model, one needs to feed a correct config.json. Remove the "false negative" part.

tianyu-l · 2025-08-07T01:19:34Z

scripts/checkpoint_conversion/checkpoint_conversion.md

+## Methods
+
+### Sanity Check (Greedy Decode)
+A quick way to sanity check if your conversion is correct is to perform greedy decoding inference on both the initial and converted checkpoints and confirm that they are the same. This method doesn't guarantee correctness but will very likely result in a fast **true negative** if the model definitions are not the same. For greedy decoding, the `generation/test_generate.py` script can be used.


This test_generate is for llama 3 only. In general we don't have such thing, and shouldn't rely on them anyways. So if you still would like to have it here, let's explicitly say something like "it's only available for llama 3, but the methodology is general".

tianyu-l · 2025-08-07T01:20:16Z

scripts/checkpoint_conversion/checkpoint_conversion.md

+### Comprehensive Check (KL Divergence)
+To ensure comprehensive end-to-end correctness we recommend using KL divergence loss to compare the logits between forward passes of both the original and converted model definitions. KL divergence quantifies the "difference" between two probability distributions. A result of zero or a very low KL divergence indicates that the model definitions are equivalent. This method is crucial as it evaluates the entire probability distribution, not just the highest probability at each step.
+
+In our `./scripts/checkpoint_conversion/example.py` this will be performing forward on dcp checkpoints loaded in `torchtitan` and safetensors checkpoints loaded in huggingface `AutoModelForCausalLM`. We additionally compare the conversions done with no permutation to double check that our permutation results in a lower kl divergence loss.


You need to provide some context on why this permutation is needed in the first place. Otherwise people will get confused why you mention it at all.

tianyu-l

Looks quite good. Had some final comments.

tianyu-l · 2025-08-08T23:58:58Z

scripts/checkpoint_conversion/example.py

+        state_dict.pop(k, None)
+
+    # Checkpoint Loading
+    logger.info(f"Loading chkpt at: {checkpoint_path}")


Suggested change

logger.info(f"Loading chkpt at: {checkpoint_path}")

logger.info(f"Loading checkpoint at: {checkpoint_path}")

tianyu-l · 2025-08-08T23:59:59Z

scripts/checkpoint_conversion/example.py

+    hf_model_path = "outputs/checkpoint/step-0-tohf"
+    hf_model_path_no_perm = "outputs/checkpoint/step-0-tohfnoperm"
+
+    # tt params
+    config_path = "torchtitan/models/llama3/train_configs/llama3_8b.toml"
+    baseline_checkpoint_path = "outputs/checkpoint/step-0-fromllama"
+    checkpoint_path = "outputs/checkpoint/step-0-fromhf"
+    checkpoint_path_no_perm = "outputs/checkpoint/step-0-fromhfnoperm"


Could you add comment: what are these checkpoints and how they are generated / downloaded?
The point is -- for any one working on a new model, they know what to do.

tianyu-l · 2025-08-09T00:00:47Z

scripts/checkpoint_conversion/README.md

now reads very well!

tianyu-l · 2025-08-09T00:01:50Z

scripts/checkpoint_conversion/README.md

+### Comprehensive Check (KL Divergence)
+To ensure comprehensive end-to-end correctness we recommend using KL divergence loss to compare the logits between forward passes of both the original and converted model definitions. KL divergence quantifies the "difference" between two probability distributions. A result of zero or a very low KL divergence indicates that the model definitions are equivalent. This method is crucial as it evaluates the entire probability distribution, not just the highest probability at each step.
+
+In our `./scripts/checkpoint_conversion/example.py` this will be performing forward on dcp checkpoints loaded in `torchtitan` and safetensors checkpoints loaded in huggingface `AutoModelForCausalLM`. To convert Llama3 between HuggingFace and `torchtitan` we had to perform a permutation on several of the attention matrices to account for difference between HuggingFace and native Llama RoPE implementations. To demonstrate how a kl divergence test can reveal subtle inaccuracies such as this, we additionally compare the kl divergence between the original and converted model with and without the permutation. The results are as follows:


Suggested change

In our `./scripts/checkpoint_conversion/example.py` this will be performing forward on dcp checkpoints loaded in `torchtitan` and safetensors checkpoints loaded in huggingface `AutoModelForCausalLM`. To convert Llama3 between HuggingFace and `torchtitan` we had to perform a permutation on several of the attention matrices to account for difference between HuggingFace and native Llama RoPE implementations. To demonstrate how a kl divergence test can reveal subtle inaccuracies such as this, we additionally compare the kl divergence between the original and converted model with and without the permutation. The results are as follows:

In our `./scripts/checkpoint_conversion/example.py` this will be performing forward on dcp checkpoints loaded in `torchtitan` and safetensors checkpoints loaded in huggingface `AutoModelForCausalLM`. To convert Llama3 between HuggingFace and `torchtitan` we had to perform a permutation on several of the attention matrices to account for difference between HuggingFace and native Llama RoPE implementations. To demonstrate how a KL divergence test can reveal subtle inaccuracies such as this, we additionally compare the KL divergence between the original and converted model with and without the permutation. The results are as follows:

tianyu-l · 2025-08-09T00:02:55Z

scripts/checkpoint_conversion/example.py

let's name it numerical_tests_example.py

wesleytruong requested review from tianyu-l, fegin, wwwjn and wconstab as code owners August 6, 2025 18:57

meta-cla bot added the CLA Signed This label is managed by the Meta Open Source bot. label Aug 6, 2025

added example for bidirectional checkpoint testing

f5f9f14

wesleytruong force-pushed the hf_conversion_tests branch from e3e1be8 to f5f9f14 Compare August 6, 2025 19:56

tianyu-l reviewed Aug 7, 2025

View reviewed changes

Made checkpoint_conversion README more clear

d495e68

tianyu-l reviewed Aug 9, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

added example for bidirectional checkpoint testing #1540

added example for bidirectional checkpoint testing #1540

Uh oh!

wesleytruong commented Aug 6, 2025

Uh oh!

tianyu-l Aug 7, 2025

Uh oh!

tianyu-l Aug 7, 2025

Uh oh!

tianyu-l Aug 7, 2025

Uh oh!

tianyu-l Aug 7, 2025

Uh oh!

tianyu-l Aug 7, 2025

Uh oh!

tianyu-l left a comment

Uh oh!

tianyu-l Aug 8, 2025

Uh oh!

tianyu-l Aug 8, 2025

Uh oh!

tianyu-l Aug 9, 2025

Uh oh!

tianyu-l Aug 9, 2025

Uh oh!

tianyu-l Aug 9, 2025

Uh oh!

Uh oh!

	logger.info(f"Loading chkpt at: {checkpoint_path}")
	logger.info(f"Loading checkpoint at: {checkpoint_path}")

added example for bidirectional checkpoint testing #1540

Are you sure you want to change the base?

added example for bidirectional checkpoint testing #1540

Uh oh!

Conversation

wesleytruong commented Aug 6, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tianyu-l left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!