A Neuro-Symbolic Legal Auditor & Formal Verification Agent
📖 Overview
JurisSim-32B v3.1 is a neuro-symbolic legislative stress-tester designed for the AMD Instinct MI300X. It translates natural language legal clauses into Z3 SMT-LIB formal constraints to identify adversarial loopholes with mathematical certainty.
🚀 Version 3.1: Dual-Agent Swarm
We have moved beyond a single model to a Multi-Agent Feedback Loop:
The Auditor (Qwen3-32B): Translates legalese into symbolic logic.
The Skeptic (Qwen2.5-7B): A dedicated linter agent that dry-runs the generated Z3 code and performs real-time self-correction on syntax errors.
Training Method: QLoRA (Quantized Low-Rank Adaptation) using bitsandbytes
2. The Fine-Tuning Pipeline & ROCm Optimizations
Training a massive 32-billion parameter model on a single MI300X GPU required severe optimization to avoid mathematical overflows and Out of Memory (OOM) crashes.
The PyTorch SDPA ROCm Bug: During initial training runs, PyTorch's default sdpa (Scaled Dot-Product Attention) implementation suffered a math overflow bug on ROCm 6.2 when handling sequences with padding tokens alongside gradient checkpointing, resulting in catastrophic NaN gradients.
The "Ultra-Stable" Workaround: Because compiling flash_attention_2 for ROCm from source is time-prohibitive, we engineered an ultra-stable configuration that perfectly maximized the 192GB VRAM without hitting the bug:
Attention: Reverted to PyTorch's native eager mathematical attention block to avoid sdpa math corruption.
Memory Compression: Reduced per_device_train_batch_size=1 but massively increased gradient_accumulation_steps=8 to maintain an effective batch size of 8.
Checkpointing: Enabled gradient_checkpointing=True to prevent the eager attention matrices from consuming all 192GB of VRAM during the backward pass.
Evaluation Safety: Enforced per_device_eval_batch_size=1 to ensure the un-checkpointed validation phase did not crash the GPU.
3. Training Results
The model successfully converged after 3 Epochs (1,713 Steps).
Final Train Loss:1.756
Final Validation Loss:1.676(Signaling excellent generalization and zero overfitting).
Token Prediction Accuracy:62.61%(Extremely high for complex Legal English → Python Z3 logic translation).
Gradient Stability: Maintained a remarkably stable grad_norm of ~0.3 to ~0.8 throughout the run.
(The Agentic Execution loop and Frontend UI are currently being finalized. Instructions for launching the UI will be placed here shortly).
🛠️ Hugging Face Space
The interactive demo of JurisSim will be deployed as a Hugging Face Space for the hackathon judging. The link will be provided upon completion of the UI.
