feat: restore native Gemma4 support

[aoshen02]

Summary

• restore Gemma4 as a native Megatron model plugin under vime_plugins/models/
• register direct HF → Megatron conversion through hf_to_megatron/_LOADERS
• register direct Megatron → HF conversion through megatron_to_hf/convert_to_hf
• keep checkpoint export and rollout weight update on the unchanged HfWeightIteratorDirect path

Architecture

This follows Slime's native architecture, not external Megatron-Bridge:

• model/spec/provider: restored from Slime's last native Gemma4 implementation before THUDM/slime#2251, translated only from slime_plugins to vime_plugins
• conversion: integrated into Slime's current direct loader/converter dispatch pattern
• shared callers: no changes to model_provider.py, hf_checkpoint_saver.py, or hf_weight_iterator_base.py
• environment: no Docker changes, Bridge patch, Bridge installation, or megatron.bridge runtime import

The only Vime-local adaptations are required by the current dependency/runtime contract: Transformers 5.15 heterogeneous Gemma4 config fields, expert partition metadata for direct HF loading, and exposing layer_scalar to the generic direct iterator.

Fixes #388.

Validation

Gemma4-26B-A4B-it on two 8×H200 nodes, TP2 / PP2 / EP4 / ETP2:

• HF → Megatron → HF completed on all 16 ranks
• tensors: 657 expected / 657 produced
• missing/extra keys: 0 / 0
• shape/dtype mismatches: 0
• bitwise exact: 657 / 657
• expert-boundary mismatches: 0

Additional checks:

• focused committed tests: 29 passed
• restored historical Slime native-model tests: 42 passed
• pre-commit: passed
• megatron.bridge availability in validation image: None

[gemini-code-assist]

Code Review

This pull request integrates support for MoE Gemma4 models using Megatron-Bridge. It updates the Docker setup to install Megatron-Bridge, adds a bridge implementation for Gemma4, integrates weight loading and saving paths, and includes corresponding unit tests. The review feedback suggests several robustness and code quality improvements: adding defensive checks for None inputs in is_gemma4_bridge_model, replacing an assert statement with a ValueError for safer input validation, and simplifying the weight replacement logic by using a standard Python list instead of a custom iterator class.

[gemini-code-assist]

If config_or_path is None (which can happen if args.hf_checkpoint is not set or is None in certain configurations), is_gemma4_bridge_model will raise an AttributeError when trying to access text_config.model_type.

Adding a defensive check for None at the beginning of the function and using getattr safely prevents potential crashes during model provider initialization.

Suggested change

	def is_gemma4_bridge_model(config_or_path: PretrainedConfig | str | Path) -> bool:
	config = (
	AutoConfig.from_pretrained(config_or_path, trust_remote_code=True)
	if isinstance(config_or_path, (str, Path))
	else config_or_path
	)
	text_config = _text_config(config)
	return text_config.model_type == "gemma4_text" and bool(getattr(text_config, "enable_moe_block", False))
	def is_gemma4_bridge_model(config_or_path: PretrainedConfig | str | Path | None) -> bool:
	if config_or_path is None:
	return False
	config = (
	AutoConfig.from_pretrained(config_or_path, trust_remote_code=True)
	if isinstance(config_or_path, (str, Path))
	else config_or_path
	)
	text_config = _text_config(config)
	if text_config is None:
	return False
	return getattr(text_config, "model_type", None) == "gemma4_text" and bool(getattr(text_config, "enable_moe_block", False))

[gemini-code-assist]

Using assert statements for application logic or input validation is an anti-pattern in Python because assertions can be globally disabled when Python is run with the -O (optimize) flag.

It is safer and more robust to raise a ValueError or TypeError instead.

Suggested change

	assert config is None, "Gemma4 builds its Megatron config from the HuggingFace checkpoint"
	if config is not None:
	raise ValueError("Gemma4 builds its Megatron config from the HuggingFace checkpoint; custom config is not supported.")

[gemini-code-assist]

The custom _MappedTasks class adds unnecessary complexity and overhead. Since the number of conversion tasks corresponds to the model's parameters (typically a few hundreds), we can simplify this by returning a standard Python list directly.

This improves readability, eliminates the need for the custom helper class, and avoids potential compatibility issues with libraries expecting a standard sequence/list.

Suggested change

	def _replace_task_weights(tasks, local_weights):
	def replace(task):
	if task is None or task.param_weight is None:
	return task
	key = task.global_param_name
	if key not in local_weights:
	if not isinstance(task.param_weight, torch.nn.Parameter):
	return task
	raise KeyError(f"Megatron-Bridge conversion weight is missing: {key}")
	return dataclasses.replace(task, param_weight=local_weights[key].cuda())
	return _MappedTasks(replace, tasks)
	class _MappedTasks:
	def __init__(self, fn, tasks):
	self.fn = fn
	self.tasks = tasks
	def __len__(self):
	return len(self.tasks)
	def __iter__(self):
	return map(self.fn, self.tasks)
	def _replace_task_weights(tasks, local_weights):
	replaced = []
	for task in tasks:
	if task is None or task.param_weight is None:
	replaced.append(task)
	continue
	key = task.global_param_name
	if key not in local_weights:
	if not isinstance(task.param_weight, torch.nn.Parameter):
	replaced.append(task)
	continue
	raise KeyError(f"Megatron-Bridge conversion weight is missing: {key}")
	replaced.append(dataclasses.replace(task, param_weight=local_weights[key].cuda()))
	return replaced