Qwen3-235B-A22B OOMs with sufficient amount of VRAM

[krammnic]

We're encountering an out-of-memory (OOM) error on a 2TB machine when attempting to load Qwen3-235B-A22B.

The issue seems to be that torchtune loads the full, unsharded checkpoint with self._checkpoint_client.load_base_checkpoint(). For a bf16 model of this size, the complete weights (~3.7TB) exceed the available memory. Loading the model in shards would likely resolve this.

Has anyone else encountered this?

Originally posted by @leng-yue in #2867 (comment)

[krammnic]

@leng-yue Hey! Thanks for the pointing this out. Would you mind to share your config? I will repro on my machine today.

[SalmanMohammadi]

@krammnic I don't think we natively support the MoE Qwen3 models? I believe there's a PR up for it though #2820.

[krammnic]

@SalmanMohammadi Oops. I'm curios how they have even managed to load the checkpoint.

[krammnic]

@leng-yue Have you used #2820? Would you mind to clarify your setup little bit?