Skip to content

Commit 7e5db60

Browse files
committed
refactor: address Gemini review comments
- _build_standard_jobs now accepts job_fn parameter; microscale single-file fallback reuses it instead of duplicating the loop Signed-off-by: David Zheng <dqzheng1996@gmail.com>
1 parent 577087d commit 7e5db60

File tree

1 file changed

+12
-26
lines changed

1 file changed

+12
-26
lines changed

src/llmcompressor/entrypoints/model_free/__init__.py

Lines changed: 12 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -121,22 +121,15 @@ def _build_standard_jobs(
121121
ignore: Iterable[str],
122122
device: torch.device,
123123
converter: Converter | None,
124+
job_fn: callable = process_file,
124125
) -> list[tuple]:
125-
"""Build one job per safetensors file for non-microscale schemes."""
126+
"""Build one job per safetensors file using the given processing function."""
126127
jobs = []
127128
for file_path, resolved_path in model_files.items():
128129
if file_path.endswith("safetensors"):
129130
save_path = Path(save_directory) / file_path
130131
jobs.append(
131-
(
132-
process_file,
133-
resolved_path,
134-
save_path,
135-
scheme,
136-
ignore,
137-
device,
138-
converter,
139-
)
132+
(job_fn, resolved_path, save_path, scheme, ignore, device, converter)
140133
)
141134
return jobs
142135

@@ -163,22 +156,15 @@ def _build_microscale_jobs(
163156

164157
if index_file is None:
165158
# Single-file model (no index.json) — use standard microscale path
166-
jobs = []
167-
for file_path, resolved_path in model_files.items():
168-
if file_path.endswith("safetensors"):
169-
save_path = Path(save_directory) / file_path
170-
jobs.append(
171-
(
172-
process_file_microscale_scheme,
173-
resolved_path,
174-
save_path,
175-
scheme,
176-
ignore,
177-
device,
178-
converter,
179-
)
180-
)
181-
return jobs
159+
return _build_standard_jobs(
160+
model_files,
161+
save_directory,
162+
scheme,
163+
ignore,
164+
device,
165+
converter,
166+
job_fn=process_file_microscale_scheme,
167+
)
182168

183169
# Read weight map to determine cross-shard fused weight groupings
184170
with open(index_file, "r") as f:

0 commit comments

Comments
 (0)