feat: add memory pre-checks and optimize PassGPT training for large wordlists

Training previously loaded entire wordlists into RAM and tokenized all at once, causing OOM on large files like rockyou.txt. This adds memory estimation, lazy dataset loading, and training optimizations. - Add _get_available_memory_mb() for cross-platform RAM detection - Add _estimate_training_memory_mb() to predict peak usage before loading - Replace bulk tokenization with LazyPasswordDataset (file offset index + on-the-fly tokenization) - Add --max-lines flag to limit training to first N lines - Add --memory-limit flag to auto-tune --max-lines based on available RAM - Enable gradient checkpointing and gradient accumulation (steps=4) - Enable fp16 on CUDA devices Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-07-30 23:50:31 -07:00 · 2026-02-18 10:47:44 -05:00
parent c3c4d9da60
commit b4bfba5fed
2 changed files with 426 additions and 63 deletions
@@ -8,8 +8,10 @@ from __future__ import annotations

 import argparse
 import os
+import subprocess
 import sys

+
 # Disable HuggingFace telemetry before any HF imports
 os.environ["HF_HUB_DISABLE_TELEMETRY"] = "1"

@@ -30,6 +32,78 @@ def _configure_mps() -> None:
    os.environ.setdefault("PYTORCH_MPS_LOW_WATERMARK_RATIO", "0.3")


+def _get_available_memory_mb() -> int | None:
+    """Return available system RAM in MB, or None if detection fails."""
+    try:
+        if sys.platform == "linux":
+            with open("/proc/meminfo") as f:
+                for line in f:
+                    if line.startswith("MemAvailable:"):
+                        return int(line.split()[1]) // 1024
+            return None
+        elif sys.platform == "darwin":
+            # macOS: try os.sysconf first, fall back to sysctl
+            try:
+                page_size = os.sysconf("SC_PAGE_SIZE")
+                avail_pages = os.sysconf("SC_AVPHYS_PAGES")
+                if page_size > 0 and avail_pages > 0:
+                    return (page_size * avail_pages) // (1024 * 1024)
+            except (ValueError, OSError):
+                pass
+            # Fallback: use sysctl for total memory (not available, but better than nothing)
+            try:
+                result = subprocess.run(
+                    ["sysctl", "-n", "hw.memsize"],
+                    capture_output=True,
+                    text=True,
+                    timeout=5,
+                )
+                if result.returncode == 0:
+                    return int(result.stdout.strip()) // (1024 * 1024)
+            except (subprocess.TimeoutExpired, FileNotFoundError, ValueError):
+                pass
+            return None
+        else:
+            return None
+    except Exception:
+        return None
+
+
+def _count_lines(filepath: str) -> int:
+    """Count non-empty lines in a file without loading it into memory."""
+    count = 0
+    with open(filepath, encoding="utf-8", errors="replace") as f:
+        for line in f:
+            if line.strip():
+                count += 1
+    return count
+
+
+def _estimate_training_memory_mb(
+    training_file: str, max_length: int = 16, max_lines: int = 0
+) -> int:
+    """Estimate peak memory usage in MB for training on the given file.
+
+    Components:
+    - Model: ~500MB (GPT-2 small)
+    - Optimizer states: ~1000MB (2x model for AdamW momentum/variance)
+    - Dataset offset index: ~8 bytes per line
+    - Per-batch activations and tokenization buffer: ~200MB
+    """
+    num_lines = _count_lines(training_file)
+    if max_lines > 0:
+        num_lines = min(num_lines, max_lines)
+
+    model_mb = 500
+    optimizer_mb = 1000
+    # Offset index: 8 bytes per line (Python int in list)
+    index_mb = (num_lines * 8) // (1024 * 1024)
+    # Activation/buffer overhead
+    buffer_mb = 200
+
+    return model_mb + optimizer_mb + index_mb + buffer_mb
+
+
 def train(
    training_file: str,
    output_dir: str,
@@ -37,7 +111,43 @@ def train(
    epochs: int,
    batch_size: int,
    device: str | None,
+    max_lines: int = 0,
+    memory_limit: int = 0,
 ) -> None:
+    # --- Memory pre-check ---
+    if memory_limit > 0:
+        # Auto-tune max_lines to fit within memory_limit
+        estimated_base = _estimate_training_memory_mb(training_file, max_lines=1)
+        per_line_bytes = 8  # offset index cost per line
+        available_for_data = (memory_limit - estimated_base) * 1024 * 1024
+        if available_for_data > 0:
+            auto_max_lines = available_for_data // per_line_bytes
+            if max_lines == 0 or auto_max_lines < max_lines:
+                max_lines = max(1, int(auto_max_lines))
+                print(
+                    f"[*] --memory-limit {memory_limit}MB: auto-set --max-lines to {max_lines}",
+                    file=sys.stderr,
+                )
+        else:
+            print(
+                f"[!] --memory-limit {memory_limit}MB is too low for model overhead alone.",
+                file=sys.stderr,
+            )
+            sys.exit(1)
+
+    estimated = _estimate_training_memory_mb(training_file, max_lines=max_lines)
+    available = _get_available_memory_mb()
+    if available is not None and estimated > available:
+        print(
+            f"[!] Estimated memory usage ({estimated}MB) exceeds available RAM ({available}MB).",
+            file=sys.stderr,
+        )
+        print(
+            "[!] Use --max-lines to limit wordlist size or --memory-limit to auto-tune.",
+            file=sys.stderr,
+        )
+        sys.exit(1)
+
    if device == "mps" or device is None:
        _configure_mps()

@@ -56,40 +166,68 @@ def train(
    tokenizer = RobertaTokenizerFast.from_pretrained(base_model)
    model = GPT2LMHeadModel.from_pretrained(base_model).to(device)  # type: ignore[arg-type]

-    print(f"[*] Reading training file: {training_file}", file=sys.stderr)
-    with open(training_file, encoding="utf-8", errors="replace") as f:
-        passwords = [line.strip() for line in f if line.strip()]
-    print(f"[*] Loaded {len(passwords)} passwords", file=sys.stderr)
-
-    print("[*] Tokenizing passwords...", file=sys.stderr)
-    max_length = model.config.n_positions if hasattr(model.config, "n_positions") else 16
-    encodings = tokenizer(
-        passwords,
-        truncation=True,
-        padding="max_length",
-        max_length=max_length,
-        return_tensors="pt",
+    max_length = (
+        model.config.n_positions if hasattr(model.config, "n_positions") else 16
    )

-    class PasswordDataset(torch.utils.data.Dataset):  # type: ignore[type-arg]
-        def __init__(self, encodings):
-            self.input_ids = encodings["input_ids"]
-            self.attention_mask = encodings["attention_mask"]
+    # Enable gradient checkpointing to reduce activation memory
+    model.gradient_checkpointing_enable()

-        def __len__(self):
-            return len(self.input_ids)
+    print(f"[*] Indexing training file: {training_file}", file=sys.stderr)

-        def __getitem__(self, idx):
+    class LazyPasswordDataset(torch.utils.data.Dataset):  # type: ignore[type-arg]
+        """Dataset that indexes file byte offsets and tokenizes on-the-fly."""
+
+        def __init__(
+            self,
+            filepath: str,
+            tokenizer: object,
+            max_length: int,
+            max_lines: int = 0,
+        ):
+            self.filepath = filepath
+            self.tokenizer = tokenizer
+            self.max_length = max_length
+            self.offsets: list[int] = []
+            with open(filepath, "rb") as f:
+                while True:
+                    offset = f.tell()
+                    line = f.readline()
+                    if not line:
+                        break
+                    if line.strip():
+                        self.offsets.append(offset)
+                        if max_lines > 0 and len(self.offsets) >= max_lines:
+                            break
+
+        def __len__(self) -> int:
+            return len(self.offsets)
+
+        def __getitem__(self, idx: int) -> dict[str, object]:
+            with open(self.filepath, "rb") as f:
+                f.seek(self.offsets[idx])
+                line = f.readline().decode("utf-8", errors="replace").strip()
+            enc = self.tokenizer(  # type: ignore[operator]
+                line,
+                truncation=True,
+                padding="max_length",
+                max_length=self.max_length,
+                return_tensors="pt",
+            )
+            input_ids = enc["input_ids"].squeeze(0)
+            attention_mask = enc["attention_mask"].squeeze(0)
            return {
-                "input_ids": self.input_ids[idx],
-                "attention_mask": self.attention_mask[idx],
-                "labels": self.input_ids[idx],
+                "input_ids": input_ids,
+                "attention_mask": attention_mask,
+                "labels": input_ids,
            }

-    dataset = PasswordDataset(encodings)
+    dataset = LazyPasswordDataset(training_file, tokenizer, max_length, max_lines)
+    print(f"[*] Indexed {len(dataset)} passwords", file=sys.stderr)

    # Use CPU for training args if device is MPS (Trainer handles device placement)
    use_cpu = device not in ("cuda",)
+    use_fp16 = device == "cuda"
    training_args = TrainingArguments(
        output_dir=output_dir,
        num_train_epochs=epochs,
@@ -99,6 +237,9 @@ def train(
        use_cpu=use_cpu,
        report_to="none",
        push_to_hub=False,
+        gradient_accumulation_steps=4,
+        fp16=use_fp16,
+        gradient_checkpointing=True,
    )

    trainer = Trainer(
@@ -159,6 +300,18 @@ def main() -> None:
        default=None,
        help="Device: cuda, mps, or cpu (default: auto-detect)",
    )
+    parser.add_argument(
+        "--max-lines",
+        type=int,
+        default=0,
+        help="Limit training to the first N lines of the wordlist (default: 0, no limit)",
+    )
+    parser.add_argument(
+        "--memory-limit",
+        type=int,
+        default=0,
+        help="Memory cap in MB; auto-tunes --max-lines to fit (default: 0, no limit)",
+    )
    args = parser.parse_args()
    train(
        training_file=args.training_file,
@@ -167,6 +320,8 @@ def main() -> None:
        epochs=args.epochs,
        batch_size=args.batch_size,
        device=args.device,
+        max_lines=args.max_lines,
+        memory_limit=args.memory_limit,
    )


@@ -4,6 +4,12 @@ from unittest.mock import MagicMock, patch

 import pytest

+from hate_crack.passgpt_train import (
+    _count_lines,
+    _estimate_training_memory_mb,
+    _get_available_memory_mb,
+)
+

@pytest.fixture
 def main_module(hc_module):
@@ -13,15 +19,17 @@ def main_module(hc_module):

 class TestHcatPassGPT:
    def test_builds_correct_pipe_commands(self, main_module):
-        with patch.object(main_module, "hcatBin", "hashcat"), patch.object(
-            main_module, "hcatTuning", "--force"
-        ), patch.object(main_module, "hcatPotfilePath", ""), patch.object(
-            main_module, "hcatHashFile", "/tmp/hashes.txt", create=True
-        ), patch.object(
-            main_module, "passgptModel", "javirandor/passgpt-10characters"
-        ), patch.object(main_module, "passgptBatchSize", 1024), patch(
-            "hate_crack.main.subprocess.Popen"
-        ) as mock_popen:
+        with (
+            patch.object(main_module, "hcatBin", "hashcat"),
+            patch.object(main_module, "hcatTuning", "--force"),
+            patch.object(main_module, "hcatPotfilePath", ""),
+            patch.object(main_module, "hcatHashFile", "/tmp/hashes.txt", create=True),
+            patch.object(
+                main_module, "passgptModel", "javirandor/passgpt-10characters"
+            ),
+            patch.object(main_module, "passgptBatchSize", 1024),
+            patch("hate_crack.main.subprocess.Popen") as mock_popen,
+        ):
            mock_gen_proc = MagicMock()
            mock_gen_proc.stdout = MagicMock()
            mock_hashcat_proc = MagicMock()
@@ -50,15 +58,17 @@ class TestHcatPassGPT:
            assert "/tmp/hashes.txt" in hashcat_cmd

    def test_custom_model_and_batch_size(self, main_module):
-        with patch.object(main_module, "hcatBin", "hashcat"), patch.object(
-            main_module, "hcatTuning", "--force"
-        ), patch.object(main_module, "hcatPotfilePath", ""), patch.object(
-            main_module, "hcatHashFile", "/tmp/hashes.txt", create=True
-        ), patch.object(
-            main_module, "passgptModel", "javirandor/passgpt-10characters"
-        ), patch.object(main_module, "passgptBatchSize", 1024), patch(
-            "hate_crack.main.subprocess.Popen"
-        ) as mock_popen:
+        with (
+            patch.object(main_module, "hcatBin", "hashcat"),
+            patch.object(main_module, "hcatTuning", "--force"),
+            patch.object(main_module, "hcatPotfilePath", ""),
+            patch.object(main_module, "hcatHashFile", "/tmp/hashes.txt", create=True),
+            patch.object(
+                main_module, "passgptModel", "javirandor/passgpt-10characters"
+            ),
+            patch.object(main_module, "passgptBatchSize", 1024),
+            patch("hate_crack.main.subprocess.Popen") as mock_popen,
+        ):
            mock_gen_proc = MagicMock()
            mock_gen_proc.stdout = MagicMock()
            mock_hashcat_proc = MagicMock()
@@ -84,9 +94,12 @@ class TestHcatPassGPTTrain:
        training_file = tmp_path / "wordlist.txt"
        training_file.write_text("password123\nabc456\n")

-        with patch.object(
-            main_module, "passgptModel", "javirandor/passgpt-10characters"
-        ), patch("hate_crack.main.subprocess.Popen") as mock_popen:
+        with (
+            patch.object(
+                main_module, "passgptModel", "javirandor/passgpt-10characters"
+            ),
+            patch("hate_crack.main.subprocess.Popen") as mock_popen,
+        ):
            mock_proc = MagicMock()
            mock_proc.returncode = 0
            mock_proc.wait.return_value = None
@@ -143,9 +156,12 @@ class TestHcatPassGPTTrain:
        training_file = tmp_path / "wordlist.txt"
        training_file.write_text("test\n")

-        with patch.object(
-            main_module, "passgptModel", "javirandor/passgpt-10characters"
-        ), patch("hate_crack.main.subprocess.Popen") as mock_popen:
+        with (
+            patch.object(
+                main_module, "passgptModel", "javirandor/passgpt-10characters"
+            ),
+            patch("hate_crack.main.subprocess.Popen") as mock_popen,
+        ):
            mock_proc = MagicMock()
            mock_proc.returncode = 1
            mock_proc.wait.return_value = None
@@ -191,8 +207,9 @@ class TestPassGPTAttackHandler:

        # "1" selects default model, "" accepts default max candidates
        inputs = iter(["1", ""])
-        with patch("builtins.input", side_effect=inputs), patch(
-            "hate_crack.attacks.os.path.isdir", return_value=False
+        with (
+            patch("builtins.input", side_effect=inputs),
+            patch("hate_crack.attacks.os.path.isdir", return_value=False),
        ):
            from hate_crack.attacks import passgpt_attack

@@ -217,13 +234,15 @@ class TestPassGPTAttackHandler:

        # "2" selects the local model, "" accepts default max candidates
        inputs = iter(["2", ""])
-        with patch("builtins.input", side_effect=inputs), patch(
-            "hate_crack.attacks.os.path.isdir", return_value=True
-        ), patch("hate_crack.attacks.os.listdir", return_value=["my_model"]), patch(
-            "hate_crack.attacks.os.path.isfile", return_value=True
-        ), patch(
-            "hate_crack.attacks.os.path.isdir",
-            side_effect=lambda p: True,
+        with (
+            patch("builtins.input", side_effect=inputs),
+            patch("hate_crack.attacks.os.path.isdir", return_value=True),
+            patch("hate_crack.attacks.os.listdir", return_value=["my_model"]),
+            patch("hate_crack.attacks.os.path.isfile", return_value=True),
+            patch(
+                "hate_crack.attacks.os.path.isdir",
+                side_effect=lambda p: True,
+            ),
        ):
            from hate_crack.attacks import passgpt_attack

@@ -241,8 +260,9 @@ class TestPassGPTAttackHandler:

        # "T" for train, "" for default base model, "" for default max candidates
        inputs = iter(["T", "", ""])
-        with patch("builtins.input", side_effect=inputs), patch(
-            "hate_crack.attacks.os.path.isdir", return_value=False
+        with (
+            patch("builtins.input", side_effect=inputs),
+            patch("hate_crack.attacks.os.path.isdir", return_value=False),
        ):
            from hate_crack.attacks import passgpt_attack

@@ -261,8 +281,9 @@ class TestPassGPTAttackHandler:
        ctx.hcatPassGPTTrain.return_value = None

        inputs = iter(["T", ""])
-        with patch("builtins.input", side_effect=inputs), patch(
-            "hate_crack.attacks.os.path.isdir", return_value=False
+        with (
+            patch("builtins.input", side_effect=inputs),
+            patch("hate_crack.attacks.os.path.isdir", return_value=False),
        ):
            from hate_crack.attacks import passgpt_attack

@@ -289,8 +310,9 @@ class TestPassGPTAttackHandler:

        # "1" selects default model, "500000" for custom max candidates
        inputs = iter(["1", "500000"])
-        with patch("builtins.input", side_effect=inputs), patch(
-            "hate_crack.attacks.os.path.isdir", return_value=False
+        with (
+            patch("builtins.input", side_effect=inputs),
+            patch("hate_crack.attacks.os.path.isdir", return_value=False),
        ):
            from hate_crack.attacks import passgpt_attack

@@ -303,3 +325,189 @@ class TestPassGPTAttackHandler:
            model_name="javirandor/passgpt-10characters",
            batch_size=1024,
        )
+
+
+class TestGetAvailableMemoryMb:
+    def test_returns_int_or_none(self):
+        result = _get_available_memory_mb()
+        assert result is None or isinstance(result, int)
+
+    def test_never_crashes_on_any_platform(self):
+        # Should not raise regardless of platform
+        _get_available_memory_mb()
+
+    def test_returns_positive_when_detected(self):
+        result = _get_available_memory_mb()
+        if result is not None:
+            assert result > 0
+
+
+class TestCountLines:
+    def test_counts_non_empty_lines(self, tmp_path):
+        f = tmp_path / "test.txt"
+        f.write_text("line1\nline2\n\nline3\n")
+        assert _count_lines(str(f)) == 3
+
+    def test_empty_file(self, tmp_path):
+        f = tmp_path / "empty.txt"
+        f.write_text("")
+        assert _count_lines(str(f)) == 0
+
+
+class TestEstimateTrainingMemoryMb:
+    def test_returns_reasonable_estimate(self, tmp_path):
+        f = tmp_path / "words.txt"
+        f.write_text("password\n" * 1000)
+        estimate = _estimate_training_memory_mb(str(f))
+        # Should include at least model + optimizer overhead (~1700MB)
+        assert estimate >= 1700
+
+    def test_max_lines_reduces_estimate(self, tmp_path):
+        f = tmp_path / "words.txt"
+        f.write_text("password\n" * 100000)
+        full = _estimate_training_memory_mb(str(f))
+        limited = _estimate_training_memory_mb(str(f), max_lines=100)
+        assert limited <= full
+
+
+class TestMemoryPrecheck:
+    def test_aborts_when_insufficient(self, tmp_path):
+        f = tmp_path / "words.txt"
+        f.write_text("password\n" * 10)
+
+        with (
+            patch("hate_crack.passgpt_train._get_available_memory_mb", return_value=1),
+            patch(
+                "hate_crack.passgpt_train._estimate_training_memory_mb",
+                return_value=5000,
+            ),
+            pytest.raises(SystemExit),
+        ):
+            from hate_crack.passgpt_train import train
+
+            train(
+                training_file=str(f),
+                output_dir=str(tmp_path / "out"),
+                base_model="test",
+                epochs=1,
+                batch_size=1,
+                device="cpu",
+            )
+
+    def test_skips_when_detection_fails(self, tmp_path):
+        """When memory detection returns None, training proceeds past the pre-check."""
+        f = tmp_path / "words.txt"
+        f.write_text("password\n" * 10)
+
+        mock_tokenizer = MagicMock()
+        mock_model = MagicMock()
+        mock_model.config.n_positions = 16
+        mock_trainer = MagicMock()
+
+        with (
+            patch(
+                "hate_crack.passgpt_train._get_available_memory_mb", return_value=None
+            ),
+            patch(
+                "hate_crack.passgpt_train._estimate_training_memory_mb",
+                return_value=5000,
+            ),
+            patch("hate_crack.passgpt_train._configure_mps"),
+            patch(
+                "transformers.RobertaTokenizerFast.from_pretrained",
+                return_value=mock_tokenizer,
+            ),
+            patch(
+                "transformers.GPT2LMHeadModel.from_pretrained",
+                return_value=mock_model,
+            ),
+            patch("transformers.Trainer", return_value=mock_trainer),
+            patch("transformers.TrainingArguments"),
+        ):
+            from hate_crack.passgpt_train import train
+
+            train(
+                training_file=str(f),
+                output_dir=str(tmp_path / "out"),
+                base_model="test",
+                epochs=1,
+                batch_size=1,
+                device="cpu",
+            )
+
+        mock_trainer.train.assert_called_once()
+
+
+class TestMaxLines:
+    def test_count_lines_respects_limit(self, tmp_path):
+        f = tmp_path / "words.txt"
+        f.write_text("password\n" * 1000)
+        # _count_lines doesn't have a limit, but _estimate uses max_lines
+        total = _count_lines(str(f))
+        assert total == 1000
+
+    def test_estimate_uses_max_lines(self, tmp_path):
+        f = tmp_path / "words.txt"
+        f.write_text("password\n" * 10000)
+        est_full = _estimate_training_memory_mb(str(f))
+        est_limited = _estimate_training_memory_mb(str(f), max_lines=10)
+        assert est_limited <= est_full
+
+
+class TestMemoryLimitAutoTune:
+    def test_auto_tunes_max_lines(self, tmp_path, capsys):
+        f = tmp_path / "words.txt"
+        f.write_text("password\n" * 100)
+
+        mock_tokenizer = MagicMock()
+        mock_model = MagicMock()
+        mock_model.config.n_positions = 16
+        mock_trainer = MagicMock()
+
+        with (
+            patch(
+                "hate_crack.passgpt_train._get_available_memory_mb", return_value=None
+            ),
+            patch("hate_crack.passgpt_train._configure_mps"),
+            patch(
+                "transformers.RobertaTokenizerFast.from_pretrained",
+                return_value=mock_tokenizer,
+            ),
+            patch(
+                "transformers.GPT2LMHeadModel.from_pretrained",
+                return_value=mock_model,
+            ),
+            patch("transformers.Trainer", return_value=mock_trainer),
+            patch("transformers.TrainingArguments"),
+        ):
+            from hate_crack.passgpt_train import train
+
+            train(
+                training_file=str(f),
+                output_dir=str(tmp_path / "out"),
+                base_model="test",
+                epochs=1,
+                batch_size=1,
+                device="cpu",
+                memory_limit=2000,
+            )
+
+        captured = capsys.readouterr()
+        assert "--memory-limit 2000MB: auto-set --max-lines" in captured.err
+
+    def test_memory_limit_too_low_exits(self, tmp_path):
+        f = tmp_path / "words.txt"
+        f.write_text("password\n" * 10)
+
+        with pytest.raises(SystemExit):
+            from hate_crack.passgpt_train import train
+
+            train(
+                training_file=str(f),
+                output_dir=str(tmp_path / "out"),
+                base_model="test",
+                epochs=1,
+                batch_size=1,
+                device="cpu",
+                memory_limit=1,  # 1MB - way too low
+            )