fix(models): pin Transformers device_map to GPU 0 when auto sharding breaks sampling

[Aafiya-H]

Summary

• Fix CUDA device-side assert failures during sampled generation when TransformersEngine loads with device_map="auto" on multi-GPU hosts (e.g. vision1 with 4× A40). Accelerate can shard small models across GPUs and produce invalid logits, which triggers torch.multinomial to fail with a device-side assert.
• After load, probe logits with a short forward pass; if they are invalid, reload the model on GPU 0 (device_map={"": 0}) before any sampled generation, since the GPU context cannot be recovered in-process after an assert.
• Refactor model loading into reusable helpers (_assemble_model_kwargs, _load_model_weights, _drop_model_weights) to support pin-and-reload without reloading the tokenizer.
• Retry generation/streaming once if a late CUDA assert still occurs; raise a clear restart hint if reload fails after the context is poisoned.
• Harden unload() so CUDA cache cleanup does not crash after prior CUDA errors.
• Ensure examples/basic/calculator_agent.py closes the agent in a finally block before model.unload() so demo/interactive runs exit cleanly on generation failure.

Root cause

On multi-GPU nodes, device_map="auto" produced collapsed/NaN logits (e.g. max ~1.6 instead of ~26). Forward passes looked fine in some cases, but sampled generation failed with:

probability tensor contains either inf, nan or element < 0

Pinning the full model to GPU 0 restored healthy logits and generation.

Test plan

• python examples/basic/calculator_agent.py — 5/5 tests pass on vision1 (multi-GPU, no CUDA_VISIBLE_DEVICES)
• Verified warning appears when auto sharding is detected: model reloads on GPU 0 before sampling
• Confirmed explicit .to("cuda") / device_map={"": 0} works; device_map="auto" alone fails on this host
• Run on a single-GPU machine to confirm device_map="auto" is unchanged when logits probe passes