RoPE config: enforcing dim

[Andrei-Aksionov]

Hey there 👋

I’m still reviewing the code. The current step is configurations.

Congratulations on the implementation:

• looks neat
• doesn't require a lot of steps to set-up a config
• and, more importantly, doesn't have too many layers of abstractions 🎉.

---

I've noticed this TODO comment:

torchtitan/torchtitan/models/llama3/__init__.py

Lines 106 to 113 in fe80b63

	rope=RoPE.Config(
	# TODO: find better ways to enforce dim = decoder dim // n_heads, for all models
	dim=dim // n_heads,
	max_seq_len=131072,
	theta=500000,
	backend="complex",
	scaling="llama",
	),

I would like to contribute, but have a couple of questions to the author of the comment, @tianyu-l:

1. What did you mean by enforcing? It's already tied to $$head\_dim = hidden\_dim/n\_heads$$.
2. We can add a check that the parameter is accurate. Since rope.dim is not always equal to head_dim (like in Qwen 3.5), this check should be done in the model's class. In Llama3Model __init__ we have access to the full config:
   

   torchtitan/torchtitan/models/llama3/model.py

   Lines 130 to 135 in fe80b63

   	def __init__(self, config: Config):
   	super().__init__(config)
   	self.enable_weight_tying = config.enable_weight_tying
   	if self.enable_weight_tying:
   	self.tok_embeddings.weight = self.output.weight

so we can add a check there:

    def __init__(self, config: Config):
        super().__init__(config)
        ...
        head_dim = config.dim // config.layer.attention.n_heads
        if config.rope.dim != head_dim:
            raise ValueError("Write something scary about RoPE dim and head dim mismatch ...")

and similar for other models.

But, does it make any sense? 🤔
Feels like over-engineering.

3. Or, perhaps, enforcing meant to not provide rope.dim in the config at all and calculate it in the model's class?
   IMHO, it will add too much magic.

---

Overall, seems like the current approach is fine.
Nevertheless, curious to hear what idea was behind that TODO comment.

[tianyu-l]

I think we can remove this TODO. cc @fegin

But in general just fyi, the config system is still evolving.

[Andrei-Aksionov]

But in general just fyi, the config system is still evolving.

There is no doubt about it.
But even in the current state it's quite good.
Plus some small things, like

torchtitan/torchtitan/config/configurable.py

Lines 158 to 171 in 30ca37a

	if kwargs_in_config:
	# All kwargs are config fields: validate and absorb into clone.
	for key, value in kwargs.items():
	if hasattr(self, key):
	existing = getattr(self, key)
	if isinstance(existing, torch.Tensor):
	mismatch = not torch.equal(existing, value)
	else:
	mismatch = existing != value
	if mismatch:
	raise ValueError(
	f"{type(self).__name__}.build() conflict for "
	f"'{key}': config has {existing!r} "
	f"but got {value!r}"

that should make a config to be a source of truth, so it makes it harder to override values somewhere down the pipeline and thus prevent head-scratching issues.
Or @dataclass(..., slots=True) that should prevent from silent errors with misspelled assignments.
No code is perfect, I get it, but this attention to details makes a lot of difference in the end.