Fix non-QEP Runner quantization for large sharded MPS checkpoints

[y-vectorfield]

Summary

When running OneComp with qep=False on MPS, quantization can be unstable for large sharded checkpoints.

This change adjusts Runner's non-QEP quantization path for MPS so that model loading and quantization avoid the unstable device placement pattern. The QEP path is unchanged.

Validation

• Manually reran the qep=False MPS checkpoint-loading flow that previously failed near the start of shard loading with the leaked semaphore warning
• Confirmed it passes the previous failure point without emitting the warning

[y-vectorfield]

@FKKimura さん、修正が完了しました。再レビューをお願いします。

[aki916f]

ご対応ありがとうございます。
qep=Falseかつbatch_sizeを指定した時に呼ばれる onecomp/runner_methods/chunked_quantization.py の中のrun_chunked_quantizationの中でも同様の箇所があるので良ければ変更いただきたいです。

OneCompression/onecomp/runner_methods/chunked_quantization.py

Lines 76 to 79 in 81ffb99

	# Load model
	model = model_config.load_model()
	tokenizer = model_config.load_tokenizer()
	input_device = next(model.parameters()).device

[aki916f]

テストの追加もありがとうございます。

model.to("mps") が呼ばれることは確認できていますが、実際にモデルが MPS に移動したことまでは確認できていないようにみえます。また、現在のモックでは parameters() がCPU上の param を返し続けるため、.to("mps") 後に取得される input_device もCPUのままになっています。.to("mps") が呼ばれた後は parameters() が返すパラメーターの device も mps になるようなモックを用意し、その結果 prepare_calibration_dataset() に mps が渡されることまで確認するのはいかがでしょうか。

例えば、テスト関数内で以下のように .to() の呼び出しに合わせて mock 上の device を更新できるはずです

current_device = torch.device("cpu")

def mock_to(device):
    nonlocal current_device
    current_device = torch.device(device)
    return model

def mock_parameters():
    param = MagicMock()
    param.device = current_device
    return iter([param])

model.to.side_effect = mock_to
model.parameters.side_effect = mock_parameters

その上で、

runner.prepare_calibration_dataset.assert_called_once_with(
    torch.device("mps"),
    model=model,
)

をパス条件にすると、.to("mps") が呼ばれたか否かだけでなく、その後の処理でも mps が使われていることを保証できるかと思います。

[y-vectorfield]

修正しました。これで漏れていたテスト内容も網羅されたと思います。

[aki916f]

早速ご対応いただき、ありがとうございます！見たところ問題なさそうなのでそれぞれのパスで動作確認取れたら承認させていただきます。

[aki916f]

すみません、utils/perplexity.pyなどほかの箇所でも呼ばれていました。。
https://github.com/search?q=repo%3AFujitsuResearch%2FOneCompression+model_config.load_model%28&type=code

今後新規にload_modelを呼ぶときにも毎回対応する必要が無くなることを考えると、大元のModelConfig.load_model() 内部を修正するのはいかがでしょうか。。

OneCompression/onecomp/model_config.py

Lines 84 to 140 in 81ffb99

	def load_model(self, device_map=None):
	"""Load the model
	Tries ``AutoModelForCausalLM`` first. If the model is a
	Vision-Language Model (e.g. Qwen3-VL) that is not registered
	with ``AutoModelForCausalLM``, falls back to
	``AutoModelForImageTextToText``.
	Args:
	device_map (str or None):
	Override the device placement for this load.
	If ``None`` (default), ``self.device`` is used.
	"""
	effective_device = device_map if device_map is not None else self.device
	kwargs = dict(
	dtype=self.dtype if self.dtype == "auto" else getattr(torch, self.dtype),
	device_map=effective_device,
	)
	config = self.load_config()
	qcfg = getattr(config, "quantization_config", None)
	if isinstance(qcfg, dict) and qcfg.get("quant_method") == "mxfp4":
	from transformers import Mxfp4Config
	from .utils.mxfp4_compat import patch_mxfp4_flat_blocks
	# Some MXFP4 checkpoints (e.g. Tokyotech GPT-OSS-Swallow) store
	# packed blocks in a flattened 3-D layout that transformers'
	# dequantizer rejects; normalize them on the fly.
	patch_mxfp4_flat_blocks(self.logger)
	kwargs["quantization_config"] = Mxfp4Config(dequantize=True)
	# MXFP4 dequantization targets bfloat16, and GPT-OSS (the only
	# MXFP4 architecture) degrades in float16.
	if kwargs["dtype"] != torch.bfloat16:
	self.logger.warning(
	"MXFP4 model detected; overriding dtype from %s to bfloat16.",
	kwargs["dtype"],
	)
	kwargs["dtype"] = torch.bfloat16
	self.logger.info("MXFP4 model detected; loading with dequantization enabled.")
	try:
	model = AutoModelForCausalLM.from_pretrained(self.get_model_id_or_path(), **kwargs)
	except ValueError as e:
	_vlm_hints = (
	"Unrecognized configuration class",
	"Unrecognized model",
	"is not supported",
	)
	if not _HAS_VLM_AUTO or not any(h in str(e) for h in _vlm_hints):
	raise
	self.logger.info("AutoModelForCausalLM failed; trying AutoModelForImageTextToText.")
	model = _AutoVLM.from_pretrained(self.get_model_id_or_path(), **kwargs)
	model.eval()
	self.logger.info("Model loaded with dtype=%s", next(model.parameters()).dtype)
	return model

[y-vectorfield]

これでmodel_configのload_model自体の修正、テストを追加できたと思います。

[y-vectorfield]

@aki916f さん、何度もレビューいただきありがとう御座いました。

[aki916f]

@y-vectorfield
修正ありがとうございます。batch_size指定の場合だとmps版のtorchがfloat64に対応していない関係で、ここでエラーが出ますね。。(これはbaseの時点で存在していたバグです。)

OneCompression/onecomp/runner_methods/chunked_quantization.py

Line 218 in 81ffb99

matrix_x = matrix_x.reshape(-1, matrix_x.shape[-1]).to(torch.float64)

このPRのスコープ外のように思うので、暫定対応として、お手数ですがmps指定かつbatch_size指定の場合はRunner.check()内でValueErrorで[mps指定の場合はbatch_size指定しないで動かす]ようメッセージを追加いただけますか？このPRの対応としてはそこで十分かと思います。

[y-vectorfield]

@y-vectorfield 修正ありがとうございます。batch_size指定の場合だとmps版のtorchがfloat64に対応していない関係で、ここでエラーが出ますね。。(これはbaseの時点で存在していたバグです。)

OneCompression/onecomp/runner_methods/chunked_quantization.py

Line 218 in 81ffb99

matrix_x = matrix_x.reshape(-1, matrix_x.shape[-1]).to(torch.float64)

このPRのスコープ外のように思うので、暫定対応として、お手数ですがmps指定かつbatch_size指定の場合はRunner.check()内でValueErrorで[mps指定の場合はbatch_size指定しないで動かす]ようメッセージを追加いただけますか？このPRの対応としてはそこで十分かと思います。

修正を追加しておきました。

[aki916f]

LGTM