[compat] compat moe_router_load_balancing_type mcore>=0.16

[Jintao-Huang]

No description provided.

[gemini-code-assist]

Code Review

This pull request adds support for multiple MoE router load balancing strategies and their corresponding auxiliary loss coefficients, updating both the documentation and implementation. The review feedback highlights two issues: a potential TypeError in base.py when the load balancing type is None, and an incorrect type hint for moe_aux_loss_coeff in megatron_args.py.

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]

If config.moe_router_load_balancing_type is None (which is its default value), load_balancing_type will remain None. This will cause a TypeError: argument of type 'NoneType' is not iterable when evaluating 'aux_loss' in load_balancing_type on line 143. We should handle the None case safely.

Suggested change

	load_balancing_type = config.moe_router_load_balancing_type
	if isinstance(load_balancing_type, str):
	load_balancing_type = [load_balancing_type]
	load_balancing_type = config.moe_router_load_balancing_type
	if load_balancing_type is None:
	load_balancing_type = []
	elif isinstance(load_balancing_type, str):
	load_balancing_type = [load_balancing_type]

[gemini-code-assist]

The type hint Union[List[float]] is equivalent to List[float], which conflicts with the default value of 0. (a float). Since moe_aux_loss_coeff can be either a single float or a list of floats, the type hint should be updated to Union[float, List[float]].

Suggested change

	moe_aux_loss_coeff: Union[List[float]] = 0.
	moe_aux_loss_coeff: Union[float, List[float]] = 0.