Make ModelDims.from_hf_config robust to explicit head_dim

[hamishivi]

Replacement for #1731, rebased onto current main and pushed from hamishivi/open-instruct.

Summary

• Honor an explicit head_dim from HF configs (e.g. composite/VLM models) instead of always deriving it from hidden_size // num_attention_heads, which fails when hidden_size is not divisible by num_attention_heads.
• Relax the __post_init__ assertion to require a positive head_dim.

Test plan

• Existing ModelDims tests.

GPU_TESTS=bypass

Made with Cursor

[gemini-code-assist]

Code Review

This pull request updates ModelDims.from_hf_config in open_instruct/utils.py to robustly handle configurations with an explicit head_dim attribute, falling back to calculating it from hidden_size and num_attention_heads only when it is absent. Additionally, the assertion in __post_init__ is updated to verify that head_dim is positive. The reviewer suggested explicitly casting head_dim to an integer when retrieved from the configuration to prevent potential type mismatches.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

To ensure that head_dim is always an integer (even if retrieved as a float or string from a custom Hugging Face config), it is safer to explicitly cast it to int when it is not None. This prevents potential type mismatch issues or float propagation in downstream parameter and FLOP calculations.

Suggested change

	head_dim = getattr(config, "head_dim", None)
	if head_dim is None:
	assert hidden_size % config.num_attention_heads == 0, (
	"hidden_size must be divisible by num_attention_heads"
	)
	head_dim = hidden_size // config.num_attention_heads
	head_dim = getattr(config, "head_dim", None)
	if head_dim is None:
	assert hidden_size % config.num_attention_heads == 0, (
	"hidden_size must be divisible by num_attention_heads"
	)
	head_dim = hidden_size // config.num_attention_heads
	else:
	head_dim = int(head_dim)