feat(phyai): refactor quantization methods.#42
Conversation
|
Caution Review failedThe pull request is closed. ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: ⛔ Files ignored due to path filters (2)
📒 Files selected for processing (30)
📝 WalkthroughWalkthroughThis PR adds a semantic-to-physical quantization system for PhyAI, including ChangesQuantization System
Estimated code review effort: 4 (Complex) | ~60 minutes Humanizer Skill Definitions and Misc Config
Estimated code review effort: 2 (Simple) | ~15 minutes Sequence Diagram(s)sequenceDiagram
participant Entry as ModelEntry.setup
participant Active as load_quant_plan
participant Plan as QuantPlan
participant Linear as LinearBase
participant Materialize as materialize
Entry->>Active: load_quant_plan(checkpoint_dir)
Active->>Plan: build_quant_plan(ConfigSources)
Entry->>Linear: construct model inside use_quant_plan(plan)
Linear->>Plan: resolve(prefix, module_cls)
Plan-->>Linear: QuantScheme or None
Linear->>Materialize: materialize(scheme, sm)
Materialize-->>Linear: WeightSpec
Possibly related PRs
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![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 introduces a comprehensive quantization subsystem for PhyAI, featuring a layered design that separates semantic quantization decisions (QuantScheme) from physical memory layouts and kernel selections (WeightSpec). It adds automatic quantization plan loading from checkpoints (supporting HF fp8, compressed-tensors, and NVIDIA ModelOpt formats), integrates this loading into model setup entry points, and adds extensive documentation in both English and Chinese. The review feedback highlights several robustness improvements to prevent runtime errors, such as handling None values for the SM architecture in CPU environments, and safely parsing JSON configurations where expected lists might be strings or top-level elements might not be dictionaries.
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.
| if w.dtype is QDType.FP8_E4M3: | ||
| return Fp8Spec(granularity=w.granularity, block_shape=w.block_shape) | ||
| if w.dtype is QDType.NVFP4: | ||
| layout = "128x4" if sm >= _NVFP4_128X4_MIN_SM else "linear" |
There was a problem hiding this comment.
If sm_arch() returns None (e.g., on non-CUDA or CPU-only test environments), sm will be None. Comparing None >= 100 will raise a TypeError. We should handle None gracefully by defaulting sm to 0 or checking if it is not None.
| layout = "128x4" if sm >= _NVFP4_128X4_MIN_SM else "linear" | |
| layout = "128x4" if (sm or 0) >= _NVFP4_128X4_MIN_SM else "linear" |
| ignored = cfg.get("ignored_layers") or cfg.get("modules_to_not_convert") or [] | ||
| rules = tuple(Rule(Matcher("name", name), None) for name in ignored) |
There was a problem hiding this comment.
If ignored_layers or modules_to_not_convert is configured as a single string instead of a list of strings (e.g., "lm_head"), iterating over ignored will iterate over its characters (e.g., "l", "m", etc.), creating incorrect rules. We should check if ignored is a string and wrap it in a list to prevent this.
| ignored = cfg.get("ignored_layers") or cfg.get("modules_to_not_convert") or [] | |
| rules = tuple(Rule(Matcher("name", name), None) for name in ignored) | |
| ignored = cfg.get("ignored_layers") or cfg.get("modules_to_not_convert") or [] | |
| if isinstance(ignored, str): | |
| ignored = [ignored] | |
| rules = tuple(Rule(Matcher("name", name), None) for name in ignored) |
| excluded = block.get("exclude_modules") or block.get("ignore") or [] | ||
| rules = tuple(Rule(Matcher("name", name), None) for name in excluded) |
There was a problem hiding this comment.
If exclude_modules or ignore is configured as a single string instead of a list of strings, iterating over excluded will iterate over its characters, creating incorrect rules. We should check if excluded is a string and wrap it in a list.
| excluded = block.get("exclude_modules") or block.get("ignore") or [] | |
| rules = tuple(Rule(Matcher("name", name), None) for name in excluded) | |
| excluded = block.get("exclude_modules") or block.get("ignore") or [] | |
| if isinstance(excluded, str): | |
| excluded = [excluded] | |
| rules = tuple(Rule(Matcher("name", name), None) for name in excluded) |
| for name in cfg.get("ignore") or []: | ||
| rules.append(Rule(_target_matcher(name), None)) |
There was a problem hiding this comment.
If ignore is configured as a single string instead of a list of strings, iterating over it will iterate over its characters. We should check if it is a string and wrap it in a list.
| for name in cfg.get("ignore") or []: | |
| rules.append(Rule(_target_matcher(name), None)) | |
| ignore = cfg.get("ignore") or [] | |
| if isinstance(ignore, str): | |
| ignore = [ignore] | |
| for name in ignore: | |
| rules.append(Rule(_target_matcher(name), None)) |
| for target in group.get("targets") or []: | ||
| rules.append(Rule(_target_matcher(target), scheme)) |
There was a problem hiding this comment.
If targets is configured as a single string instead of a list of strings, iterating over it will iterate over its characters, creating incorrect rules. We should check if targets is a string and wrap it in a list.
| for target in group.get("targets") or []: | |
| rules.append(Rule(_target_matcher(target), scheme)) | |
| targets = group.get("targets") or [] | |
| if isinstance(targets, str): | |
| targets = [targets] | |
| for target in targets: | |
| rules.append(Rule(_target_matcher(target), scheme)) |
| hf_quant_config = raw.get("quantization_config") or raw.get( | ||
| "compression_config" | ||
| ) |
There was a problem hiding this comment.
If config.json contains a non-dictionary top-level element (e.g., a list or a string), calling raw.get will raise an AttributeError. We should check if raw is a dictionary before calling .get().
| hf_quant_config = raw.get("quantization_config") or raw.get( | |
| "compression_config" | |
| ) | |
| if isinstance(raw, dict): | |
| hf_quant_config = raw.get("quantization_config") or raw.get( | |
| "compression_config" | |
| ) |
| standalone = None | ||
| standalone_path = folder / "hf_quant_config.json" | ||
| if standalone_path.is_file(): | ||
| with standalone_path.open("r", encoding="utf-8") as f: | ||
| standalone = json.load(f) |
There was a problem hiding this comment.
If hf_quant_config.json contains a non-dictionary top-level element, standalone will not be a dictionary, and calling src.standalone.get in _quant_block will raise an AttributeError. We should check if the loaded JSON is a dictionary before assigning it to standalone.
| standalone = None | |
| standalone_path = folder / "hf_quant_config.json" | |
| if standalone_path.is_file(): | |
| with standalone_path.open("r", encoding="utf-8") as f: | |
| standalone = json.load(f) | |
| standalone = None | |
| standalone_path = folder / "hf_quant_config.json" | |
| if standalone_path.is_file(): | |
| with standalone_path.open("r", encoding="utf-8") as f: | |
| raw_standalone = json.load(f) | |
| if isinstance(raw_standalone, dict): | |
| standalone = raw_standalone |
1 participant
Summary by CodeRabbit
New Features
Documentation