@tuned-tensor/cli 0.25.0 → 0.26.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -26,6 +26,7 @@ from model_contract import (
26
26
  assert_certified_base_model_revision,
27
27
  assert_certified_model_config,
28
28
  chat_template_kwargs,
29
+ cuda_model_load_device_map,
29
30
  parse_lora_target_modules,
30
31
  uncertified_load_kwargs,
31
32
  )
@@ -216,13 +217,16 @@ def create_model_and_tokenizer(model_source: str) -> tuple[Any, Any, torch.dtype
216
217
  if hp("model_loader", "causal_lm") == "image_text_to_text"
217
218
  else AutoModelForCausalLM
218
219
  )
220
+ load_device_map = cuda_model_load_device_map(torch)
219
221
  model = model_class.from_pretrained(
220
222
  model_source,
221
223
  **source_kwargs,
222
224
  **uncertified_load_kwargs(str(hp("base_model", CERTIFIED_BASE_MODEL))),
223
225
  dtype=dtype,
224
- device_map={"": torch.cuda.current_device()},
226
+ device_map=load_device_map,
225
227
  )
228
+ if load_device_map is None:
229
+ model = model.to("cuda")
226
230
  assert_certified_model_config(
227
231
  model.config,
228
232
  expected_model_id=str(hp("base_model", CERTIFIED_BASE_MODEL)),