attacklm 0.9.0__tar.gz → 0.9.2__tar.gz

{attacklm-0.9.0 → attacklm-0.9.2}/PKG-INFO

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: attacklm
-Version: 0.9.0
-Summary: QLoRA-fine-tuned Qwen2.5-Coder for offensive/defensive security and AI red-teaming
+Version: 0.9.2
+Summary: A high-performance fine-tuning pipeline (QLoRA, GaLore, Q-GaLore, Spectrum, PiSSA, DeepSpeed) for creating MITRE ATT&CK-grounded security AI assistants
 Project-URL: Homepage, https://github.com/Veedubin/AttackLM
 Project-URL: Repository, https://github.com/Veedubin/AttackLM
 Project-URL: Issues, https://github.com/Veedubin/AttackLM/issues
@@ -26,19 +26,16 @@ Provides-Extra: all
 Requires-Dist: accelerate==1.14.0; extra == 'all'
 Requires-Dist: bitsandbytes==0.49.2; extra == 'all'
 Requires-Dist: datasets==5.0.0; extra == 'all'
+Requires-Dist: deepspeed>=0.16; extra == 'all'
 Requires-Dist: gguf; extra == 'all'
 Requires-Dist: huggingface-hub==1.18.0; extra == 'all'
-Requires-Dist: ipython; extra == 'all'
 Requires-Dist: kernels<0.14,>=0.12; extra == 'all'
 Requires-Dist: llama-cpp-python; extra == 'all'
-Requires-Dist: mypy; extra == 'all'
+Requires-Dist: lomo-optim>=0.1; extra == 'all'
 Requires-Dist: peft==0.19.1; extra == 'all'
-Requires-Dist: pytest; extra == 'all'
 Requires-Dist: pyyaml; extra == 'all'
 Requires-Dist: requests; extra == 'all'
-Requires-Dist: ruff; extra == 'all'
 Requires-Dist: safetensors==0.7.0; extra == 'all'
-Requires-Dist: scipy>=1.11; extra == 'all'
 Requires-Dist: tokenizers==0.22.2; extra == 'all'
 Requires-Dist: torch==2.12.0; extra == 'all'
 Requires-Dist: torchvision==0.27.0; extra == 'all'
@@ -49,10 +46,12 @@ Provides-Extra: all-cuda
 Requires-Dist: accelerate==1.14.0; extra == 'all-cuda'
 Requires-Dist: bitsandbytes==0.49.2; extra == 'all-cuda'
 Requires-Dist: datasets==5.0.0; extra == 'all-cuda'
+Requires-Dist: deepspeed>=0.16; extra == 'all-cuda'
 Requires-Dist: gguf; extra == 'all-cuda'
 Requires-Dist: huggingface-hub==1.18.0; extra == 'all-cuda'
 Requires-Dist: kernels<0.14,>=0.12; extra == 'all-cuda'
 Requires-Dist: llama-cpp-python; extra == 'all-cuda'
+Requires-Dist: lomo-optim>=0.1; extra == 'all-cuda'
 Requires-Dist: peft==0.19.1; extra == 'all-cuda'
 Requires-Dist: pyyaml; extra == 'all-cuda'
 Requires-Dist: requests; extra == 'all-cuda'
@@ -67,10 +66,12 @@ Provides-Extra: all-rocm
 Requires-Dist: accelerate==1.14.0; extra == 'all-rocm'
 Requires-Dist: bitsandbytes==0.49.2; extra == 'all-rocm'
 Requires-Dist: datasets==5.0.0; extra == 'all-rocm'
+Requires-Dist: deepspeed>=0.16; extra == 'all-rocm'
 Requires-Dist: gguf; extra == 'all-rocm'
 Requires-Dist: huggingface-hub==1.18.0; extra == 'all-rocm'
 Requires-Dist: kernels<0.14,>=0.12; extra == 'all-rocm'
 Requires-Dist: llama-cpp-python; extra == 'all-rocm'
+Requires-Dist: lomo-optim>=0.1; extra == 'all-rocm'
 Requires-Dist: peft==0.19.1; extra == 'all-rocm'
 Requires-Dist: pyyaml; extra == 'all-rocm'
 Requires-Dist: requests; extra == 'all-rocm'
@@ -149,8 +150,10 @@ Provides-Extra: train
 Requires-Dist: accelerate==1.14.0; extra == 'train'
 Requires-Dist: bitsandbytes==0.49.2; extra == 'train'
 Requires-Dist: datasets==5.0.0; extra == 'train'
+Requires-Dist: deepspeed>=0.16; extra == 'train'
 Requires-Dist: huggingface-hub==1.18.0; extra == 'train'
 Requires-Dist: kernels<0.14,>=0.12; extra == 'train'
+Requires-Dist: lomo-optim>=0.1; extra == 'train'
 Requires-Dist: peft==0.19.1; extra == 'train'
 Requires-Dist: safetensors==0.7.0; extra == 'train'
 Requires-Dist: tokenizers==0.22.2; extra == 'train'
@@ -161,8 +164,10 @@ Requires-Dist: trl==1.5.1; extra == 'train'
 Provides-Extra: train-base
 Requires-Dist: accelerate==1.14.0; extra == 'train-base'
 Requires-Dist: datasets==5.0.0; extra == 'train-base'
+Requires-Dist: deepspeed>=0.16; extra == 'train-base'
 Requires-Dist: huggingface-hub==1.18.0; extra == 'train-base'
 Requires-Dist: kernels<0.14,>=0.12; extra == 'train-base'
+Requires-Dist: lomo-optim>=0.1; extra == 'train-base'
 Requires-Dist: peft==0.19.1; extra == 'train-base'
 Requires-Dist: safetensors==0.7.0; extra == 'train-base'
 Requires-Dist: tokenizers==0.22.2; extra == 'train-base'
@@ -172,8 +177,10 @@ Provides-Extra: train-cuda
 Requires-Dist: accelerate==1.14.0; extra == 'train-cuda'
 Requires-Dist: bitsandbytes==0.49.2; extra == 'train-cuda'
 Requires-Dist: datasets==5.0.0; extra == 'train-cuda'
+Requires-Dist: deepspeed>=0.16; extra == 'train-cuda'
 Requires-Dist: huggingface-hub==1.18.0; extra == 'train-cuda'
 Requires-Dist: kernels<0.14,>=0.12; extra == 'train-cuda'
+Requires-Dist: lomo-optim>=0.1; extra == 'train-cuda'
 Requires-Dist: peft==0.19.1; extra == 'train-cuda'
 Requires-Dist: safetensors==0.7.0; extra == 'train-cuda'
 Requires-Dist: tokenizers==0.22.2; extra == 'train-cuda'
@@ -185,8 +192,10 @@ Provides-Extra: train-rocm
 Requires-Dist: accelerate==1.14.0; extra == 'train-rocm'
 Requires-Dist: bitsandbytes==0.49.2; extra == 'train-rocm'
 Requires-Dist: datasets==5.0.0; extra == 'train-rocm'
+Requires-Dist: deepspeed>=0.16; extra == 'train-rocm'
 Requires-Dist: huggingface-hub==1.18.0; extra == 'train-rocm'
 Requires-Dist: kernels<0.14,>=0.12; extra == 'train-rocm'
+Requires-Dist: lomo-optim>=0.1; extra == 'train-rocm'
 Requires-Dist: peft==0.19.1; extra == 'train-rocm'
 Requires-Dist: safetensors==0.7.0; extra == 'train-rocm'
 Requires-Dist: tokenizers==0.22.2; extra == 'train-rocm'

{attacklm-0.9.0 → attacklm-0.9.2}/pyproject.toml

@@ -11,7 +11,7 @@ build-backend = "hatchling.build"
 [project]
 name = "attacklm"
 dynamic = ["version"]
-description = "QLoRA-fine-tuned Qwen2.5-Coder for offensive/defensive security and AI red-teaming"
+description = "A high-performance fine-tuning pipeline (QLoRA, GaLore, Q-GaLore, Spectrum, PiSSA, DeepSpeed) for creating MITRE ATT&CK-grounded security AI assistants"
 readme = "README.md"
 license = "MIT"
 requires-python = ">=3.10"
@@ -105,6 +105,8 @@ train-base = [
     "huggingface-hub==1.18.0",
     "tokenizers==0.22.2",
     "safetensors==0.7.0",
+    "deepspeed>=0.16",
+    "lomo-optim>=0.1",
 ]
 infer-base = [
     "transformers==5.12.1",
@@ -217,7 +219,7 @@ eval = [
 # at install time. Resolving `attacklm[all]` on a build/CI server without
 # that index URL would fail.
 all = [
-    "attacklm[train-cuda,extract,convert,dev]",
+    "attacklm[train-cuda,extract,convert]",
 ]
 all-cuda = ["attacklm[train-cuda,extract,convert]"]
 all-rocm = ["attacklm[train-rocm,extract,convert]"]

{attacklm-0.9.0 → attacklm-0.9.2}/scripts/train_template.py

@@ -565,6 +565,15 @@ Examples:
             "Takes precedence over --deepspeed-offload if both are set."
         ),
     )
+    parser.add_argument(
+        "--force-deepspeed-cuda",
+        action="store_true",
+        default=False,
+        help=(
+            "Skip DeepSpeed CUDA version check even when major versions differ. "
+            "Only use if you know your CUDA toolkit is ABI-compatible with PyTorch's CUDA."
+        ),
+    )
 
     # ---- torch.compile ----
     parser.add_argument(
@@ -3180,6 +3189,57 @@ def main() -> None:
     # With CPU offload, a 16GB GPU + 64GB RAM can train 40B+ parameter models.
     deepspeed_config_path = None
     if args.use_deepspeed:
+        # DeepSpeed compiles CUDA extensions at runtime and checks that the
+        # system CUDA version exactly matches what PyTorch was compiled with.
+        # CUDA guarantees ABI compatibility within a major version (e.g. 13.0
+        # and 13.3 are compatible), so we only skip the check when the major
+        # versions match.  If they differ, we warn and let the user force it.
+        import re, subprocess
+
+        _torch_cuda = getattr(torch.version, "cuda", None) or ""
+        _sys_cuda = ""
+        try:
+            _nvcc = subprocess.run(
+                ["nvcc", "--version"], capture_output=True, text=True, timeout=5
+            )
+            _m = re.search(r"release\s+(\d+\.\d+)", _nvcc.stdout)
+            if _m:
+                _sys_cuda = _m.group(1)
+        except Exception:
+            pass
+
+        if _torch_cuda and _sys_cuda:
+            _torch_major = _torch_cuda.split(".")[0]
+            _sys_major = _sys_cuda.split(".")[0]
+            if _torch_major == _sys_major:
+                os.environ.setdefault("DS_SKIP_CUDA_CHECK", "1")
+                console.print(
+                    f"  [dim]DeepSpeed:[/dim] CUDA {_sys_cuda} (system) ≈ "
+                    f"CUDA {_torch_cuda} (torch) — same major version, "
+                    f"skipping strict version check"
+                )
+            else:
+                console.print(
+                    f"[yellow]WARNING:[/yellow] System CUDA {_sys_cuda} != "
+                    f"torch CUDA {_torch_cuda} (different major versions). "
+                    f"DeepSpeed may fail to compile extensions."
+                )
+                console.print(
+                    "  If you know these are ABI-compatible, set "
+                    "DS_SKIP_CUDA_CHECK=1 in your environment."
+                )
+                console.print("  Or pass --force-deepspeed-cuda to skip the check.")
+                if not args.force_deepspeed_cuda:
+                    console.print(
+                        "[red]ERROR:[/red] Refusing to proceed with CUDA version "
+                        "mismatch. Use --force-deepspeed-cuda to override."
+                    )
+                    sys.exit(1)
+                os.environ["DS_SKIP_CUDA_CHECK"] = "1"
+        else:
+            # Can't detect versions — skip the check and hope for the best
+            os.environ.setdefault("DS_SKIP_CUDA_CHECK", "1")
+
         # Verify deepspeed package is installed
         try:
             import deepspeed  # noqa: F401

attacklm-0.9.2/src/attacklm/__version__.py

@@ -0,0 +1 @@
+__version__ = "0.9.2"

attacklm-0.9.0/src/attacklm/__version__.py

@@ -1 +0,0 @@
-__version__ = "0.9.0"