refactor olmoe

[Jintao-Huang]

No description provided.

[gemini-code-assist]

Code Review

This pull request refactors the OLMoE model implementation, simplifying OLMoESelfAttention, OLMoEBridge, and OLMoELoader by inheriting from and leveraging base classes. It updates the QK layernorm logic to apply RMSNorm over the full Q/K channels and correctly handle tensor parallel (TP) replication. The review feedback suggests adding defensive None checks for self.q_norm and self.k_norm when configuring their parameters and applying them to the query/key tensors to prevent potential runtime errors if QK layernorm is disabled.

[gemini-code-assist]

If submodules.q_layernorm or submodules.k_layernorm is None (for example, if QK layernorm is disabled or not configured in a custom setup), self.q_norm or self.k_norm will be initialized as None. Accessing .parameters() on a None object will raise an AttributeError. Adding a defensive check to ensure the norm is not None before iterating over its parameters prevents potential runtime crashes.

Suggested change

	for norm in (self.q_norm, self.k_norm):
	for p in norm.parameters():
	p.sequence_parallel = False
	p.average_gradients_across_tp_domain = True
	for norm in (self.q_norm, self.k_norm):
	if norm is not None:
	for p in norm.parameters():
	p.sequence_parallel = False
	p.average_gradients_across_tp_domain = True

[gemini-code-assist]

Similarly, if self.q_norm or self.k_norm is None, calling them directly will raise a TypeError. Adding defensive None checks before applying the norms ensures robustness.

Suggested change

	query = self.q_norm(query)
	key = self.k_norm(key)
	if self.q_norm is not None:
	query = self.q_norm(query)
	if self.k_norm is not None:
	key = self.k_norm(key)