VRAM Settings

Documentation Index

Fetch the complete documentation index at: https://mintlify.com/pranavkrishnasuresh/chemAgent/llms.txt

Use this file to discover all available pages before exploring further.

​

Overview

​

Configuration File

# Flag to avoid running LLaSmol with <15GB VRAM (MIN REQUIREMENT)
LOW_VRAM = True

​

VRAM Modes

​

Low VRAM Mode (Default)

LOW_VRAM = True

• LlaSMol model is NOT loaded (chem_tools.py:115-119)
• answer_chemistry_query tool will raise a RuntimeError if called
• Agent relies entirely on OpenAI GPT-4o for chemistry queries
• Suitable for systems with less than 15GB VRAM

RuntimeError: answer_chemistry_query tool cannot be used with LOW_VRAM enabled.

"LlaSmol model unused. Low VRAM enabled."

​

High VRAM Mode (Cluster/GPU)

LOW_VRAM = False

• Minimum 15GB VRAM
• CUDA-enabled GPU
• PyTorch with CUDA support

• LlaSMol model is loaded into GPU memory
• answer_chemistry_query tool becomes available
• Model uses bfloat16 precision for memory efficiency
• Automatic device mapping with device_map="auto"

​

Implementation Details

​

Conditional Loading

# Tool to use LlaSmol to answer prompts related to Chemistry
# Won't initialize with low VRAM
if not LOW_VRAM:
    from LLM4Chem.generation import LlaSMolGeneration
    generator = LlaSMolGeneration("osunlp/LlaSMol-Mistral-7B", device="cuda")
else:
    generator = None

​

Runtime Checks

if LOW_VRAM:
    llasmol_response.model_response = "LlaSmol model unused. Low VRAM enabled."
    raise RuntimeError(
        "answer_chemistry_query tool cannot be used with LOW_VRAM enabled."
    )

​

Model Memory Usage

​

Memory Optimizations

1. bfloat16 Precision (model.py:38, 45)
   model = AutoModelForCausalLM.from_pretrained(
       base_model,
       torch_dtype=torch.bfloat16,
       device_map="auto",
   )
   
2. PEFT/LoRA Loading (model.py:42-46)
   model = PeftModelForCausalLM.from_pretrained(
       model,
       model_name,
       torch_dtype=torch.bfloat16,
   )
   
3. Model Merging (model.py:50)
   model = model.merge_and_unload()
   
4. Torch Compilation (model.py:58-59)
   if torch.__version__ >= "2" and sys.platform != "win32":
       model = torch.compile(model)
   

​

Device Selection

def get_device():
    if torch.cuda.is_available():
        device = "cuda"
    else:
        device = "cpu"
    return device

​

Configuration for Different Environments

​

Local Development (Low VRAM)

LOW_VRAM = True

• Laptops with consumer GPUs (less than 15GB VRAM)
• Development machines with limited GPU memory
• Testing agent logic without model inference

​

Cluster/Production (High VRAM)

LOW_VRAM = False

• NVIDIA A100 (40GB/80GB)
• NVIDIA V100 (16GB/32GB)
• NVIDIA RTX 3090 (24GB)
• Cloud GPU instances with ≥15GB VRAM

​

Troubleshooting

​

Out of Memory Errors

1. Verify VRAM availability:
   nvidia-smi
   
2. Check available memory:
   import torch
   print(f"Available VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB")
   
3. Set LOW_VRAM = True if VRAM < 15GB

​

Model Not Loading

# Check CUDA availability
import torch
print(f"CUDA available: {torch.cuda.is_available()}")
print(f"CUDA version: {torch.version.cuda}")

​

Related Configuration

• plan_execute_agent/config.py:2 (active flag)
• LLM4Chem/config.py:2 (legacy, not actively used)

​

Next Steps