[v1] refactor registry plugin structure and params#10641
Conversation
|
@frozenleaves @sunyi0505 @Kuangdd01 @codemayq @hiyouga It's a large pr to refactor plugin and params and influence many parts. Would you mind having a check? Please feel free to comment and I will try to fix issues. |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the plugin architecture and configuration parsing in LlamaFactory (v1), introducing typed dataclasses for PEFT, quantization, and distributed training configurations to replace unstructured dictionaries. It simplifies the BasePlugin and BaseKernel interfaces, explicitly registers model kernels, and promotes distributed parameters (like context parallel size) directly to TrainingArguments. The review feedback highlights several robustness issues, including potential AttributeError and TypeError crashes if configurations are mixed between dictionaries and dataclasses, fragile device type checks, a potential crash when sorting mixed-type keys in plugin parameter parsing, and missing fallback handling for null compute_dtype values.
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.
| """Returns the unique identifier for the kernel.""" | ||
| return cls._kernel_id | ||
| def check_deps(cls) -> None: | ||
| supported = cls._device if isinstance(cls._device, list) else [cls._device] |
There was a problem hiding this comment.
Using isinstance(cls._device, list) is fragile because subclasses might define _device as a tuple (e.g., (DeviceType.CUDA, DeviceType.NPU)). If a tuple is used, supported will become a single-element list containing the tuple, which will cause dependency checks to fail. Using isinstance(cls._device, (list, tuple)) is more robust.
| supported = cls._device if isinstance(cls._device, list) else [cls._device] | |
| supported = cls._device if isinstance(cls._device, (list, tuple)) else [cls._device] |
| if unknown: | ||
| raise ValueError( | ||
| f"Unknown params for {cls.__name__}.{params_cls.__name__}: {sorted(unknown)}. " | ||
| f"Expected: {sorted(known)}" |
There was a problem hiding this comment.
If unknown contains mixed types (e.g., strings and integers/booleans from YAML keys), calling sorted(unknown) will raise a TypeError in Python 3. Converting all keys to strings before sorting prevents this crash.
| f"Expected: {sorted(known)}" | |
| f"Unknown params for {cls.__name__}.{params_cls.__name__}: {sorted(str(k) for k in unknown)}. " |
| if dist_config is None or getattr(dist_config, "name", None) != "deepspeed": | ||
| return False | ||
|
|
||
| config_file = dist_config.get("config_file") |
There was a problem hiding this comment.
If dist_config is passed as a dataclass or custom object (like DeepSpeedParams) instead of a dict/PluginConfig, calling .get() will raise an AttributeError. Using getattr or checking for hasattr(dist_config, 'get') makes it robust to both dict-like and object-like configurations.
| config_file = dist_config.get("config_file") | |
| config_file = dist_config.get("config_file") if hasattr(dist_config, "get") else getattr(dist_config, "config_file", None) |
| def setup_deepspeed_zero3_model_loading(): | ||
| """Enable ZeRO-3-aware model loading for the registered backend config.""" | ||
| dist_config = _registered_dist_config | ||
| config_file = dist_config.get("config_file") if dist_config is not None else None |
There was a problem hiding this comment.
If dist_config is passed as a dataclass or custom object (like DeepSpeedParams) instead of a dict/PluginConfig, calling .get() will raise an AttributeError. Using getattr or checking for hasattr(dist_config, 'get') makes it robust to both dict-like and object-like configurations.
| config_file = dist_config.get("config_file") if dist_config is not None else None | |
| config_file = (dist_config.get("config_file") if hasattr(dist_config, "get") else getattr(dist_config, "config_file", None)) if dist_config is not None else None |
| if isinstance(self.compute_dtype, str): | ||
| dtype = getattr(torch, self.compute_dtype, None) | ||
| if not isinstance(dtype, torch.dtype): | ||
| raise ValueError(f"compute_dtype={self.compute_dtype!r} is not a torch dtype name.") | ||
| self.compute_dtype = dtype | ||
| elif not isinstance(self.compute_dtype, torch.dtype): | ||
| raise TypeError(f"compute_dtype must be str or torch.dtype, got {type(self.compute_dtype).__name__}.") |
There was a problem hiding this comment.
If self.compute_dtype is explicitly set to None (e.g., compute_dtype: null in YAML), __post_init__ will raise a TypeError. Handling None by defaulting to torch.float16 improves usability and backward compatibility.
| if isinstance(self.compute_dtype, str): | |
| dtype = getattr(torch, self.compute_dtype, None) | |
| if not isinstance(dtype, torch.dtype): | |
| raise ValueError(f"compute_dtype={self.compute_dtype!r} is not a torch dtype name.") | |
| self.compute_dtype = dtype | |
| elif not isinstance(self.compute_dtype, torch.dtype): | |
| raise TypeError(f"compute_dtype must be str or torch.dtype, got {type(self.compute_dtype).__name__}.") | |
| if self.compute_dtype is None: | |
| self.compute_dtype = torch.float16 | |
| if isinstance(self.compute_dtype, str): | |
| dtype = getattr(torch, self.compute_dtype, None) | |
| if not isinstance(dtype, torch.dtype): | |
| raise ValueError(f"compute_dtype={self.compute_dtype!r} is not a torch dtype name.") | |
| self.compute_dtype = dtype | |
| elif not isinstance(self.compute_dtype, torch.dtype): | |
| raise TypeError(f"compute_dtype must be str or torch.dtype, got {type(self.compute_dtype).__name__}.") |
| if export_config.get("name") == "lora": | ||
| adapters = export_config.get("adapter_name_or_path") | ||
| else: | ||
| if raw_config.name != "lora": |
There was a problem hiding this comment.
If raw_config is a standard dictionary (e.g., when called programmatically or in custom scripts), raw_config.name will raise an AttributeError. Using getattr(raw_config, 'name', None) is safer and more robust.
| if raw_config.name != "lora": | |
| if getattr(raw_config, "name", None) != "lora": |
| @@ -93,9 +93,12 @@ def __init__( | |||
| dist_name = self.args.dist_config.name if self.args.dist_config is not None else None | |||
There was a problem hiding this comment.
If self.args.dist_config is a standard dictionary, self.args.dist_config.name will raise an AttributeError. Using getattr(self.args.dist_config, 'name', None) is safer and more robust.
| dist_name = self.args.dist_config.name if self.args.dist_config is not None else None | |
| dist_name = getattr(self.args.dist_config, "name", None) |
| from ..plugins.trainer_plugins.distributed.interface import DistributedPlugin | ||
|
|
||
| self.model = DistributedPlugin(self.args.dist_config.name)( | ||
| self.model = DistributedPlugin(self.args.dist_config.name).shard_model( |
There was a problem hiding this comment.
If self.args.dist_config is a standard dictionary, self.args.dist_config.name will raise an AttributeError. Using getattr(self.args.dist_config, 'name', None) is safer and more robust.
| self.model = DistributedPlugin(self.args.dist_config.name).shard_model( | |
| self.model = DistributedPlugin(getattr(self.args.dist_config, "name", None)).shard_model( |
| if self.args.dist_config is not None and self.args.dist_config.name in ("deepspeed", "fsdp2"): | ||
| from ..plugins.trainer_plugins.distributed.hub import DistributedPlugin | ||
| from ..plugins.trainer_plugins.distributed.interface import DistributedPlugin | ||
|
|
||
| DistributedPlugin(self.args.dist_config.name).save_model( |
There was a problem hiding this comment.
If self.args.dist_config is a standard dictionary, self.args.dist_config.name will raise an AttributeError. Using getattr(self.args.dist_config, 'name', None) is safer and more robust.
| if self.args.dist_config is not None and self.args.dist_config.name in ("deepspeed", "fsdp2"): | |
| from ..plugins.trainer_plugins.distributed.hub import DistributedPlugin | |
| from ..plugins.trainer_plugins.distributed.interface import DistributedPlugin | |
| DistributedPlugin(self.args.dist_config.name).save_model( | |
| if self.args.dist_config is not None and getattr(self.args.dist_config, "name", None) in ("deepspeed", "fsdp2"): | |
| from ..plugins.trainer_plugins.distributed.interface import DistributedPlugin | |
| DistributedPlugin(getattr(self.args.dist_config, "name", None)).save_model( |
1 participant
What does this PR do?
The previous plugin registration mechanism managed both plugin implementations and their parameter types, which tightly coupled plugin discovery, configuration parsing, and invocation. Some plugin families also maintained separate registries or hubs.
This PR unifies the v1 plugin registration mechanism and moves parameter parsing to the actual plugin entry points.
Main Changes
1. Unified plugin registration
BasePluginnow maintains an independent registry for each plugin family, preventing registry data from being shared across PEFT, Kernel, Batching, and Distributed plugins.2. Parse parameters at plugin entry points
Plugin parameters are explicitly parsed through the corresponding plugin family’s
parse_params()method where the plugin is actually invoked.For example, LoRA parameters are parsed at the PEFT entry point:
This keeps the registry focused on plugin implementations, while each plugin family owns its parameter parsing and each entry point controls when parsing occurs. Quantization and distributed plugins follow the same pattern.
3. Kernel plugin updates
KernelPlugin.4. Batching plugin updates
5. Distributed plugin updates
DistributedInterface.dist_config.6. Distributed parallelism parameters
Common parallel topology settings are no longer stored in backend-specific
dist_config. They are managed directly byTrainingArgumentsand passed to the distributed interface:dp_sizecp_sizecp_modemp_replicate_sizemp_shard_sizedist_timeoutThis separates general parallelism topology from backend-specific configuration.
Sequence Parallel keeps its existing directory and module structure, with
cp_sizenow passed explicitly at the call site.7. Configuration and example updates
Before submitting