From b4bfba5fed98251863298bf2f9bb0f115062609c Mon Sep 17 00:00:00 2001 From: Justin Bollinger Date: Wed, 18 Feb 2026 10:47:44 -0500 Subject: [PATCH] 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 --- hate_crack/passgpt_train.py | 203 ++++++++++++++++++++++--- tests/test_passgpt_attack.py | 286 ++++++++++++++++++++++++++++++----- 2 files changed, 426 insertions(+), 63 deletions(-) diff --git a/hate_crack/passgpt_train.py b/hate_crack/passgpt_train.py index fd8f60b..8892484 100644 --- a/hate_crack/passgpt_train.py +++ b/hate_crack/passgpt_train.py @@ -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, ) diff --git a/tests/test_passgpt_attack.py b/tests/test_passgpt_attack.py index 245f397..63fe8a6 100644 --- a/tests/test_passgpt_attack.py +++ b/tests/test_passgpt_attack.py @@ -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 + )