Add missing RoBERTa converter - #2960
Conversation
Split out of PR keras-team#2929 per review feedback, isolating RoBERTa so its numerical verification can be reviewed independently.
convert_roberta.py (the Hugging Face converter now wired into from_preset()) already handles config translation and weight mapping, so this script no longer needs to hand-port a raw fairseq checkpoint layer by layer. It now loads directly via RobertaBackbone/RobertaTokenizer.from_preset(hf://...), matching the pattern used by convert_bert_sentence_transformer_checkpoints.py and convert_smollm3_checkpoints.py. Verified end-to-end against FacebookAI/roberta-base: KerasHub output matches the HF reference model, and the resulting preset loads correctly.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request registers a new Hugging Face converter for RoBERTa, allowing direct loading of RoBERTa presets from Hugging Face, and refactors the checkpoint conversion script to leverage this new converter. Feedback on the changes includes: restoring leading underscores when accessing private attention dense layers to prevent runtime AttributeErrors, using a tiny random model in unit tests to ensure fast execution, and masking out padding positions during numerical verification to accurately compute the logits difference on content tokens.
| loader.port_weight( | ||
| keras_variable=encoder_layer._self_attention_layer.query_dense.kernel, | ||
| hf_weight_key=f"{hf_prefix}.attention.self.query.weight", | ||
| hook_fn=transpose_and_reshape, | ||
| ) | ||
| loader.port_weight( | ||
| keras_variable=encoder_layer._self_attention_layer.query_dense.bias, | ||
| hf_weight_key=f"{hf_prefix}.attention.self.query.bias", | ||
| hook_fn=lambda hf_tensor, shape: np.reshape(hf_tensor, shape), | ||
| ) | ||
| loader.port_weight( | ||
| keras_variable=encoder_layer._self_attention_layer.key_dense.kernel, | ||
| hf_weight_key=f"{hf_prefix}.attention.self.key.weight", | ||
| hook_fn=transpose_and_reshape, | ||
| ) | ||
| loader.port_weight( | ||
| keras_variable=encoder_layer._self_attention_layer.key_dense.bias, | ||
| hf_weight_key=f"{hf_prefix}.attention.self.key.bias", | ||
| hook_fn=lambda hf_tensor, shape: np.reshape(hf_tensor, shape), | ||
| ) | ||
| loader.port_weight( | ||
| keras_variable=encoder_layer._self_attention_layer.value_dense.kernel, | ||
| hf_weight_key=f"{hf_prefix}.attention.self.value.weight", | ||
| hook_fn=transpose_and_reshape, | ||
| ) | ||
| loader.port_weight( | ||
| keras_variable=encoder_layer._self_attention_layer.value_dense.bias, | ||
| hf_weight_key=f"{hf_prefix}.attention.self.value.bias", | ||
| hook_fn=lambda hf_tensor, shape: np.reshape(hf_tensor, shape), | ||
| ) | ||
| loader.port_weight( | ||
| keras_variable=encoder_layer._self_attention_layer.output_dense.kernel, | ||
| hf_weight_key=f"{hf_prefix}.attention.output.dense.weight", | ||
| hook_fn=transpose_and_reshape, | ||
| ) | ||
| loader.port_weight( | ||
| keras_variable=encoder_layer._self_attention_layer.output_dense.bias, | ||
| hf_weight_key=f"{hf_prefix}.attention.output.dense.bias", | ||
| hook_fn=lambda hf_tensor, shape: np.reshape(hf_tensor, shape), | ||
| ) |
There was a problem hiding this comment.
The attention dense layers in Keras's MultiHeadAttention (used by TransformerEncoder) are private attributes prefixed with an underscore (e.g., _query_dense, _key_dense, _value_dense, _output_dense). Accessing them as public attributes (e.g., query_dense) will raise an AttributeError at runtime when loading weights. Please restore the leading underscores to match the Keras implementation and the original checkpoint conversion script.
# Attention layers
loader.port_weight(
keras_variable=encoder_layer._self_attention_layer._query_dense.kernel,
hf_weight_key=f"{hf_prefix}.attention.self.query.weight",
hook_fn=transpose_and_reshape,
)
loader.port_weight(
keras_variable=encoder_layer._self_attention_layer._query_dense.bias,
hf_weight_key=f"{hf_prefix}.attention.self.query.bias",
hook_fn=lambda hf_tensor, shape: np.reshape(hf_tensor, shape),
)
loader.port_weight(
keras_variable=encoder_layer._self_attention_layer._key_dense.kernel,
hf_weight_key=f"{hf_prefix}.attention.self.key.weight",
hook_fn=transpose_and_reshape,
)
loader.port_weight(
keras_variable=encoder_layer._self_attention_layer._key_dense.bias,
hf_weight_key=f"{hf_prefix}.attention.self.key.bias",
hook_fn=lambda hf_tensor, shape: np.reshape(hf_tensor, shape),
)
loader.port_weight(
keras_variable=encoder_layer._self_attention_layer._value_dense.kernel,
hf_weight_key=f"{hf_prefix}.attention.self.value.weight",
hook_fn=transpose_and_reshape,
)
loader.port_weight(
keras_variable=encoder_layer._self_attention_layer._value_dense.bias,
hf_weight_key=f"{hf_prefix}.attention.self.value.bias",
hook_fn=lambda hf_tensor, shape: np.reshape(hf_tensor, shape),
)
loader.port_weight(
keras_variable=encoder_layer._self_attention_layer._output_dense.kernel,
hf_weight_key=f"{hf_prefix}.attention.output.dense.weight",
hook_fn=transpose_and_reshape,
)
loader.port_weight(
keras_variable=encoder_layer._self_attention_layer._output_dense.bias,
hf_weight_key=f"{hf_prefix}.attention.output.dense.bias",
hook_fn=lambda hf_tensor, shape: np.reshape(hf_tensor, shape),
)References
- Respect existing repo patterns and ensure robustness by avoiding fragile code that raises runtime exceptions. (link)
There was a problem hiding this comment.
query_dense/key_dense/value_dense/output_dense are public property aliases for the private attrs in Keras's MultiHeadAttention, so the current code is correct.The tests already pass without an AttributeError.
| def test_convert_tiny_preset(self): | ||
| model = RobertaTextClassifier.from_preset( | ||
| "hf://FacebookAI/roberta-base", num_classes=2 | ||
| ) |
There was a problem hiding this comment.
The test test_convert_tiny_preset is intended to verify the converter with a lightweight preset, but it currently loads the full-sized FacebookAI/roberta-base model (125M parameters, ~500MB download). This makes the unit test extremely slow and resource-intensive. Please use a tiny random model from Hugging Face, such as hf://hf-internal-testing/tiny-random-RobertaModel, to keep the test fast and lightweight.
| def test_convert_tiny_preset(self): | |
| model = RobertaTextClassifier.from_preset( | |
| "hf://FacebookAI/roberta-base", num_classes=2 | |
| ) | |
| def test_convert_tiny_preset(self): | |
| model = RobertaTextClassifier.from_preset( | |
| "hf://hf-internal-testing/tiny-random-RobertaModel", num_classes=2 | |
| ) |
References
- When a test is set up, the test inputs should be small to allow for fast testing. (link)
There was a problem hiding this comment.
the model size is under 500 mb, it should be fine. also tried the suggested preset anyways, there was an error.
| print("HF output:", hf_output[0, 0, :10]) | ||
| print("Difference:", np.mean(keras_hub_output - hf_output.detach().numpy())) |
There was a problem hiding this comment.
As noted in the PR description, RoBERTa's Hugging Face embedding reuses a constant padding_idx position row for every pad token, while KerasHub's shared PositionEmbedding layer assigns sequential positions regardless of padding. This results in a large difference (up to 3.05) at padding positions, which is expected and inert. However, np.mean(keras_hub_output - hf_output) averages over the entire sequence (including padding), printing a misleadingly high difference. We should mask out the padding positions using the padding mask to print the true difference for content tokens.
| print("HF output:", hf_output[0, 0, :10]) | |
| print("Difference:", np.mean(keras_hub_output - hf_output.detach().numpy())) | |
| padding_mask = keras_hub_inputs["padding_mask"].numpy() | |
| diff = np.abs(keras_hub_output - hf_output.detach().numpy()) | |
| non_padded_diff = diff * np.expand_dims(padding_mask, axis=-1) | |
| mean_diff = np.sum(non_padded_diff) / (np.sum(padding_mask) * keras_hub_output.shape[-1]) | |
| print("Difference (non-padding):", mean_diff) |
References
- When performing numerical verification for converted model checkpoints, validate using the mean logits difference tolerance rather than the maximum absolute difference, and ensure the comparison is accurate and robust.
There was a problem hiding this comment.
Fixed. check_output now slices to real (non-padding) tokens via padding_mask and asserts np.testing.assert_allclose(atol=1e-4, rtol=1e-4) instead of printing a whole-sequence mean.
Ahmed0830
changed the title
transformers.logging.set_verbosity_error() silences the (expected, harmless) state-dict load report that fires whenever a bare RobertaModel is loaded from an MLM-pretraining checkpoint. check_output() now asserts numerical parity (atol=1e-4) instead of only printing a diff. The comparison is restricted to real, non-padding tokens: HF assigns a constant padding_idx position embedding to every pad token, while KerasHub PositionEmbedding assigns sequential positions regardless of padding, so a whole-sequence diff is dominated by that expected, functionally-inert divergence rather than actual conversion error.
3 participants
Description of the change
Split out of #2929 per review feedback, isolating the RoBERTa HF→KerasHub converter so it can be reviewed and numerically verified independently of the other converters in that PR.
keras_hub/src/utils/transformers/convert_roberta.pyand wiresmodel_type == "roberta"intoTransformersPresetLoader.convert_roberta_test.py(test_convert_tiny_preset,test_class_detection).tools/checkpoint_conversion/convert_roberta_checkpoints.pypreviously hand-ported weights from the original fairseq RoBERTa checkpoint, layer by layer. Sinceconvert_roberta.pynow handles the Hugging Face weight mapping directly viafrom_preset(), that manual porting code is not needed, the script now loads straight from huggingface viafrom_preset(), matching the pattern inconvert_bert_sentence_transformer_checkpoints.py.Verification
Ran the rewritten
convert_roberta_checkpoints.py --preset roberta_base_enend-to-end against the liveFacebookAI/roberta-basecheckpoint and comparedlast_hidden_stateagainst the reference HF model.Colab Reference: colab
pytest --run_extra_large keras_hub/src/utils/transformers/convert_roberta_test.py— 2 passed.Checklist