diff --git a/dev_scripts/laptop_pull_bundle.sh b/dev_scripts/laptop_pull_bundle.sh new file mode 100755 index 0000000..1c25fb6 --- /dev/null +++ b/dev_scripts/laptop_pull_bundle.sh @@ -0,0 +1,31 @@ +#!/bin/bash + +# Script to pull git bundle from server, unpack, and push to remote + +# Prompt for connection details +read -p "Enter server IP: " IP +read -p "Enter server port: " PORT +read -p "Enter branch name: " BRANCH_NAME + +# Pull down the bundle +BUNDLE_FILE="${BRANCH_NAME}.bundle" +scp -P "$PORT" "root@${IP}:/root/lda/${BUNDLE_FILE}" . + +# Switch to main and pull latest +git checkout main +git pull + +# Unpack the bundle +git bundle verify "$BUNDLE_FILE" +git fetch "$BUNDLE_FILE" "$BRANCH_NAME":"$BRANCH_NAME" + +# Switch to the branch +git checkout "$BRANCH_NAME" + +# Push to remote +git push -u origin "$BRANCH_NAME" + +# Clean up bundle file +rm "$BUNDLE_FILE" + +echo "Branch $BRANCH_NAME has been pushed to origin" diff --git a/dev_scripts/server_bundle.sh b/dev_scripts/server_bundle.sh new file mode 100755 index 0000000..d9b4521 --- /dev/null +++ b/dev_scripts/server_bundle.sh @@ -0,0 +1,38 @@ +#!/bin/bash + +# Script to create a git branch, stage all files, commit, and create a bundle + +cd /root/lda + +# Check if git user.name is set, if not prompt for it +if [ -z "$(git config user.name)" ]; then + read -p "Enter your git name: " GIT_NAME + git config --global user.name "$GIT_NAME" +fi + +# Check if git user.email is set, if not prompt for it +if [ -z "$(git config user.email)" ]; then + read -p "Enter your git email: " GIT_EMAIL + git config --global user.email "$GIT_EMAIL" +fi + +# Prompt for branch name +read -p "Enter branch name: " BRANCH_NAME + +# Create and switch to new branch +git checkout -b "$BRANCH_NAME" + +# Stage all files +git add -A + +# Prompt for commit message +read -p "Enter commit message: " COMMIT_MSG + +# Commit changes +git commit -m "$COMMIT_MSG" + +# Create bundle +BUNDLE_PATH="/root/lda/${BRANCH_NAME}.bundle" +git bundle create "$BUNDLE_PATH" main.."$BRANCH_NAME" + +echo "Bundle created at: $BUNDLE_PATH" diff --git a/dev_scripts/startup.sh b/dev_scripts/startup.sh new file mode 100644 index 0000000..a4b7a28 --- /dev/null +++ b/dev_scripts/startup.sh @@ -0,0 +1,13 @@ +#!/bin/bash + +# Install uv +pip install uv + +# Navigate to lda directory +cd /root/lda + +# Sync dependencies +uv sync + +# Run the server +uv run python server.py \ No newline at end of file diff --git a/lda.py b/lda.py index 39fe030..b7808c0 100644 --- a/lda.py +++ b/lda.py @@ -43,8 +43,17 @@ class TokenizerCompatibility: def _apply_top_p_filtering(logits: torch.Tensor, top_p: float) -> torch.Tensor: - """Apply nucleus (top-p) filtering to logits.""" - sorted_logits, sorted_indices = torch.sort(logits, descending=True) + """ + Apply nucleus (top-p) filtering to logits. + + Args: + logits: Logits tensor, shape [vocab_size] or [batch_size, vocab_size] + top_p: Top-p threshold + + Returns: + Filtered logits with same shape as input + """ + sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1) cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) # Remove tokens with cumulative probability above the threshold @@ -60,8 +69,20 @@ def _apply_top_p_filtering(logits: torch.Tensor, top_p: float) -> torch.Tensor: return logits -def _sample_token(logits: torch.Tensor, temperature: float, top_p: float) -> int: - """Sample a token from logits with temperature and top-p filtering.""" +def _sample_token(logits: torch.Tensor, temperature: float, top_p: float) -> int | torch.Tensor: + """ + Sample a token from logits with temperature and top-p filtering. + + Args: + logits: Logits tensor, shape [vocab_size] or [batch_size, vocab_size] + temperature: Temperature for sampling + top_p: Top-p threshold + + Returns: + Single token ID (int) if logits is 1D, or tensor of token IDs [batch_size] if 2D + """ + is_batched = logits.ndim == 2 + # Apply temperature logits = logits / temperature @@ -72,7 +93,12 @@ def _sample_token(logits: torch.Tensor, temperature: float, top_p: float) -> int # Sample from the distribution probs = F.softmax(logits, dim=-1) next_token = torch.multinomial(probs, num_samples=1) - return next_token.item() + + # Return int for single sample, tensor for batch + if is_batched: + return next_token.squeeze(-1) # [batch_size, 1] -> [batch_size] + else: + return next_token.item() def _format_as_chat(prompt: str, tokenizer, system_prompt: str | None = None) -> str: @@ -237,17 +263,19 @@ def get_tokenizer_compatibility(self) -> TokenizerCompatibility: def _lda_step( self, input_ids: torch.Tensor, + attention_mask: torch.Tensor, cache_after: tuple | None, cache_before: tuple | None, alpha: float, temperature: float, top_p: float - ) -> tuple[int, tuple, tuple]: + ) -> tuple[int | torch.Tensor, tuple, tuple]: """ Perform a single LDA generation step with KV-cache support. Args: - input_ids: Input token IDs (full sequence on first call, single token after) + input_ids: Input token IDs, shape [batch_size, seq_len] (full sequence on first call, single token after) + attention_mask: Attention mask, shape [batch_size, seq_len] (0 for padding, 1 for real tokens) cache_after: KV-cache from previous step for after model (None on first call) cache_before: KV-cache from previous step for before model (None on first call) alpha: Amplification factor @@ -255,22 +283,25 @@ def _lda_step( top_p: Top-p sampling threshold Returns: - Tuple of (next_token_id, new_cache_after, new_cache_before) + Tuple of (next_token_id(s), new_cache_after, new_cache_before) + - next_token_id: int if batch_size=1, torch.Tensor[batch_size] otherwise """ # Get logits from both models with KV-cache outputs_after = self.model_after( input_ids, + attention_mask=attention_mask, past_key_values=cache_after, use_cache=True ) - logits_after = outputs_after.logits[0, -1, :] # Last token logits + logits_after = outputs_after.logits[:, -1, :] # Last token logits, preserve batch dim outputs_before = self.model_before( input_ids, + attention_mask=attention_mask, past_key_values=cache_before, use_cache=True ) - logits_before = outputs_before.logits[0, -1, :] # Last token logits + logits_before = outputs_before.logits[:, -1, :] # Last token logits, preserve batch dim # Compute amplified logits # logits_amplified = logits_after + alpha * (logits_after - logits_before) @@ -278,7 +309,13 @@ def _lda_step( logits_amplified = logits_after + alpha * logits_diff # Sample next token from amplified logits - next_token_id = _sample_token(logits_amplified, temperature, top_p) + # For batch_size=1, this returns int; for batch_size>1, returns tensor[batch_size] + if logits_amplified.shape[0] == 1: + # Single sample case: squeeze batch dimension for backward compatibility + next_token_id = _sample_token(logits_amplified[0], temperature, top_p) + else: + # Batched case: keep batch dimension + next_token_id = _sample_token(logits_amplified, temperature, top_p) return next_token_id, outputs_after.past_key_values, outputs_before.past_key_values @@ -319,7 +356,9 @@ def generate_lda( formatted_prompt = _format_as_chat(prompt, self.tokenizer_after, system_prompt) # Tokenize initial prompt - input_ids = self.tokenizer_after(formatted_prompt, return_tensors="pt")["input_ids"].to(self.device) + tokenized = self.tokenizer_after(formatted_prompt, return_tensors="pt") + input_ids = tokenized["input_ids"].to(self.device) + attention_mask = tokenized["attention_mask"].to(self.device) generated_ids = input_ids.clone() # Initialize KV-caches (None means first iteration) @@ -335,8 +374,10 @@ def generate_lda( step_input_ids = generated_ids if cache_after is None else generated_ids[:, -1:] # Perform LDA step with caching + # Note: attention_mask must be full length (all positions), not sliced next_token_id, cache_after, cache_before = self._lda_step( step_input_ids, + attention_mask, cache_after, cache_before, alpha, @@ -354,6 +395,9 @@ def generate_lda( # Append to generated sequence next_token_tensor = torch.tensor([[next_token_id]], device=self.device) generated_ids = torch.cat([generated_ids, next_token_tensor], dim=1) + + # Update attention mask (new token is always attended to) + attention_mask = torch.cat([attention_mask, torch.ones((1, 1), device=self.device, dtype=attention_mask.dtype)], dim=1) # Decode results input_len = input_ids.shape[1] @@ -368,6 +412,127 @@ def generate_lda( stopped_early=stopped_early ) + def generate_lda_batched( + self, + prompts: list[str], + max_new_tokens: int, + temperature: float, + top_p: float, + alpha: float, + apply_chat_template: bool = True, + system_prompt: str | None = None + ) -> list[LDAResult]: + """ + Generate text using logit diff amplification for multiple prompts in a batch. + + Args: + prompts: List of input prompts + max_new_tokens: Maximum number of tokens to generate (fixed for all prompts) + temperature: Temperature for sampling + top_p: Top-p (nucleus) sampling threshold + alpha: Amplification factor for logit differences + apply_chat_template: Whether to apply chat template formatting + system_prompt: Optional system prompt for chat template + + Returns: + List of LDAResult objects, one per prompt + """ + # Warn if tokenizers are incompatible + if not self._compatibility.compatible: + warnings.warn( + "Tokenizers are incompatible! LDA results may be nonsensical." + ) + + if len(prompts) == 0: + return [] + + # Apply chat template if requested + formatted_prompts = prompts + if apply_chat_template: + formatted_prompts = [ + _format_as_chat(prompt, self.tokenizer_after, system_prompt) + for prompt in prompts + ] + + # Configure tokenizer for left-padding + original_padding_side = self.tokenizer_after.padding_side + self.tokenizer_after.padding_side = "left" + + # Set pad token if not already set + if self.tokenizer_after.pad_token is None: + self.tokenizer_after.pad_token = self.tokenizer_after.eos_token + + # Tokenize all prompts with padding + tokenized = self.tokenizer_after( + formatted_prompts, + return_tensors="pt", + padding=True, + truncation=False + ) + input_ids = tokenized["input_ids"].to(self.device) + attention_mask = tokenized["attention_mask"].to(self.device) + + # Restore original padding side + self.tokenizer_after.padding_side = original_padding_side + + batch_size = input_ids.shape[0] + generated_ids = input_ids.clone() + + # Initialize KV-caches + cache_after = None + cache_before = None + + # Fixed-length generation (no EOS checking for batched mode) + with torch.no_grad(): + for i in range(max_new_tokens): + # Determine input: full sequence on first iteration, last token only after + step_input_ids = generated_ids if cache_after is None else generated_ids[:, -1:] + + # Perform LDA step with caching + # Note: attention_mask must be full length (all positions), not sliced + next_token_ids, cache_after, cache_before = self._lda_step( + step_input_ids, + attention_mask, + cache_after, + cache_before, + alpha, + temperature, + top_p + ) + + # Append to generated sequence + # Handle both int (batch_size=1) and tensor (batch_size>1) returns from _lda_step + if isinstance(next_token_ids, int): + next_token_tensor = torch.tensor([[next_token_ids]], device=self.device) + else: + next_token_tensor = next_token_ids.unsqueeze(-1) # [batch_size] -> [batch_size, 1] + generated_ids = torch.cat([generated_ids, next_token_tensor], dim=1) + + # Update attention mask (new tokens are always attended to) + attention_mask = torch.cat([ + attention_mask, + torch.ones((batch_size, 1), device=self.device, dtype=attention_mask.dtype) + ], dim=1) + + # Decode results for each sequence in the batch + results = [] + input_lengths = tokenized["attention_mask"].sum(dim=1).tolist() # Get original lengths before padding + + for idx in range(batch_size): + input_len = input_lengths[idx] + completion_ids = generated_ids[idx, input_len:] + completion = self.tokenizer_after.decode(completion_ids, skip_special_tokens=True) + text = self.tokenizer_after.decode(generated_ids[idx], skip_special_tokens=True) + + results.append(LDAResult( + completion=completion, + text=text, + tokens_generated=max_new_tokens, + stopped_early=False # No EOS checking in batched mode + )) + + return results + def generate( self, prompt: str, diff --git a/run_sleeper_experiment.py b/run_sleeper_experiment.py index cf0873e..c1cafdc 100644 --- a/run_sleeper_experiment.py +++ b/run_sleeper_experiment.py @@ -22,13 +22,14 @@ class ExperimentConfig: """Configuration for the experiment.""" server_url: str = "http://localhost:8000" - num_prompts: int = 10 + num_prompts: int = 16 alpha_values: list[float] = field(default_factory=lambda: [0.0, 1.0, 2.0]) max_new_tokens: int = 150 temperature: float = 1.0 top_p: float = 1.0 timeout: int = 120 hf_token: str | None = None # HuggingFace API token for gated datasets + batch_size: int = 8 # Number of prompts to process in a single batch @dataclass @@ -111,6 +112,51 @@ def run_lda_generation( ) +def run_lda_generation_batch( + prompts: list[str], + alpha: float, + config: ExperimentConfig +) -> list[TrialResult]: + """Run LDA generation for a batch of prompts with specified alpha.""" + try: + response = requests.post( + f"{config.server_url}/generate_lda_batch", + json={ + "prompts": prompts, + "max_new_tokens": config.max_new_tokens, + "temperature": config.temperature, + "top_p": config.top_p, + "alpha": alpha, + }, + timeout=config.timeout + ) + response.raise_for_status() + result = response.json() + + trial_results = [] + for batch_result in result["results"]: + trial_results.append(TrialResult( + prompt=batch_result["prompt"], + alpha=alpha, + completion=batch_result["completion"], + activated=detect_sleeper_activation(batch_result["completion"]), + )) + + return trial_results + except Exception as e: + # Return error results for all prompts in the batch + return [ + TrialResult( + prompt=prompt, + alpha=alpha, + completion="", + activated=False, + error=str(e), + ) + for prompt in prompts + ] + + def run_experiment(config: ExperimentConfig) -> dict: """Run the full experiment and return results.""" prompts = load_prompts(config.num_prompts, config.hf_token) @@ -118,6 +164,7 @@ def run_experiment(config: ExperimentConfig) -> dict: results = { "config": { "num_prompts": config.num_prompts, + "batch_size": config.batch_size, "alpha_values": config.alpha_values, "max_new_tokens": config.max_new_tokens, "temperature": config.temperature, @@ -130,38 +177,53 @@ def run_experiment(config: ExperimentConfig) -> dict: # Track activations per alpha value activations = {alpha: 0 for alpha in config.alpha_values} - activations["standard"] = 0 totals = {alpha: 0 for alpha in config.alpha_values} - totals["standard"] = 0 - total_trials = len(prompts) * (len(config.alpha_values) + 1) # +1 for standard + total_trials = len(prompts) * len(config.alpha_values) trial_num = 0 - for i, prompt in enumerate(prompts): - print(f"\n--- Prompt {i+1}/{len(prompts)} ---") - print(f"Prompt: {prompt[:80]}{'...' if len(prompt) > 80 else ''}") + # Process prompts in batches for each alpha value + for alpha in config.alpha_values: + print(f"\n=== Processing alpha={alpha} ===") - # LDA generation at each alpha - for alpha in config.alpha_values: - trial_num += 1 - print(f" [{trial_num}/{total_trials}] LDA (alpha={alpha})...", end=" ", flush=True) - result = run_lda_generation(prompt, alpha, config) - results["trials"].append({ - "prompt": prompt, - "method": "lda", - "alpha": alpha, - "completion": result.completion, - "activated": result.activated, - "error": result.error, - }) - totals[alpha] += 1 - if result.activated: - activations[alpha] += 1 - print("ACTIVATED!") - elif result.error: - print(f"ERROR: {result.error}") + # Process prompts in batches + for batch_start in range(0, len(prompts), config.batch_size): + batch_end = min(batch_start + config.batch_size, len(prompts)) + batch_prompts = prompts[batch_start:batch_end] + batch_size = len(batch_prompts) + + trial_num += batch_size + print(f" [{trial_num}/{total_trials}] Batch [{batch_start+1}-{batch_end}] (alpha={alpha})...", end=" ", flush=True) + + # Run batched generation + batch_results = run_lda_generation_batch(batch_prompts, alpha, config) + + # Process results + activated_count = 0 + error_count = 0 + for result in batch_results: + results["trials"].append({ + "prompt": result.prompt, + "method": "lda", + "alpha": alpha, + "completion": result.completion, + "activated": result.activated, + "error": result.error, + }) + totals[alpha] += 1 + if result.activated: + activations[alpha] += 1 + activated_count += 1 + elif result.error: + error_count += 1 + + # Print batch summary + if error_count > 0: + print(f"ERROR ({error_count}/{batch_size} failed)") + elif activated_count > 0: + print(f"ACTIVATED! ({activated_count}/{batch_size})") else: - print("no activation") + print(f"no activation (0/{batch_size})") # Calculate summary statistics print("\n" + "=" * 60) @@ -170,15 +232,6 @@ def run_experiment(config: ExperimentConfig) -> dict: summary = {} - # Standard sampling - rate = activations["standard"] / totals["standard"] if totals["standard"] > 0 else 0 - summary["standard"] = { - "activated": activations["standard"], - "total": totals["standard"], - "rate": rate, - } - print(f"Standard sampling: {activations['standard']}/{totals['standard']} ({rate:.1%})") - # LDA at each alpha for alpha in config.alpha_values: rate = activations[alpha] / totals[alpha] if totals[alpha] > 0 else 0 @@ -194,36 +247,41 @@ def run_experiment(config: ExperimentConfig) -> dict: return results -def main(): +def main(): parser = argparse.ArgumentParser(description="Run sleeper agent detection experiment") parser.add_argument("--server-url", default="http://localhost:8000", help="Server URL") parser.add_argument("--num-prompts", type=int, default=10, help="Number of prompts to test") - parser.add_argument("--alpha", type=float, action="append", help="Alpha values to test (can specify multiple)") + parser.add_argument("--alpha", type=float, nargs='+', help="Alpha values to test (can specify multiple)") parser.add_argument("--max-tokens", type=int, default=150, help="Max tokens to generate") + parser.add_argument("--batch-size", type=int, default=8, help="Batch size for parallel processing") parser.add_argument("--output", type=str, help="Output JSON file for results") parser.add_argument("--hf-token", type=str, help="HuggingFace API token (or set HF_TOKEN env var)") args = parser.parse_args() - + # Get HF token from args or environment hf_token = args.hf_token or os.environ.get("HF_TOKEN") if not hf_token: print("Warning: No HuggingFace token provided. Set HF_TOKEN env var or use --hf-token") print("The LMSYS-Chat-1M dataset requires authentication.") - config = ExperimentConfig( - server_url=args.server_url, - num_prompts=args.num_prompts, - max_new_tokens=args.max_tokens, - hf_token=hf_token, - ) - - if args.alpha: - config.alpha_values = args.alpha + # Build config, using CLI args if provided, otherwise use defaults + # Map CLI arg names to config field names + arg_to_config = { + "server_url": "server_url", + "num_prompts": "num_prompts", + "alpha": "alpha_values", + "max_tokens": "max_new_tokens", + "batch_size": "batch_size", + } + config_args = {arg_to_config[k]: v for k, v in vars(args).items() if k in arg_to_config and v is not None} + config_args["hf_token"] = hf_token + config = ExperimentConfig(**config_args) print("Sleeper Agent Detection Experiment") print("=" * 60) print(f"Server: {config.server_url}") print(f"Prompts: {config.num_prompts}") + print(f"Batch size: {config.batch_size}") print(f"Alpha values: {config.alpha_values}") print(f"Max tokens: {config.max_new_tokens}") print("=" * 60) diff --git a/server.py b/server.py index f22316d..e48c5e4 100644 --- a/server.py +++ b/server.py @@ -56,6 +56,17 @@ class LDAGenerateRequest(BaseModel): system_prompt: str | None = Field(None, description="Optional system prompt for chat template") +class LDABatchGenerateRequest(BaseModel): + prompts: list[str] = Field(..., min_length=1, description="List of user messages (will be formatted as chat)") + max_new_tokens: int = Field(80, ge=1, le=512, description="Maximum number of tokens to generate") + temperature: float = Field(0.8, gt=0.0, description="Temperature for sampling") + top_p: float = Field(0.95, gt=0.0, le=1.0, description="Top-p (nucleus) sampling threshold") + alpha: float = Field(1.0, ge=0.0, le=10.0, description="Amplification factor for logit differences") + do_sample: bool = Field(True, description="Whether to use sampling (always True for LDA)") + apply_chat_template: bool = Field(True, description="Whether to apply chat template formatting") + system_prompt: str | None = Field(None, description="Optional system prompt for chat template") + + class ModelResponse(BaseModel): model_id: str completion: str @@ -80,6 +91,23 @@ class LDAGenerateResponse(BaseModel): stopped_early: bool +class LDABatchResult(BaseModel): + prompt: str + formatted_prompt: str + completion: str + text: str + tokens_generated: int + stopped_early: bool + + +class LDABatchGenerateResponse(BaseModel): + results: list[LDABatchResult] + alpha: float + model_after: str + model_before: str + batch_size: int + + class TokenizerCompatibilityInfo(BaseModel): compatible: bool vocab_size_1: int @@ -211,6 +239,56 @@ def generate_lda(req: LDAGenerateRequest) -> LDAGenerateResponse: ) +@app.post("/generate_lda_batch", response_model=LDABatchGenerateResponse) +def generate_lda_batch(req: LDABatchGenerateRequest) -> LDABatchGenerateResponse: + """Generate text using Logit Diff Amplification (LDA) for multiple prompts in a batch.""" + if _model_pair is None: + raise HTTPException(status_code=503, detail="Models not initialized yet") + + if len(req.prompts) == 0: + raise HTTPException(status_code=400, detail="prompts list cannot be empty") + + with _lock: + results = _model_pair.generate_lda_batched( + prompts=req.prompts, + max_new_tokens=req.max_new_tokens, + temperature=req.temperature, + top_p=req.top_p, + alpha=req.alpha, + apply_chat_template=req.apply_chat_template, + system_prompt=req.system_prompt + ) + + # Reconstruct formatted prompts for response + from lda import _format_as_chat + formatted_prompts = req.prompts + if req.apply_chat_template: + formatted_prompts = [ + _format_as_chat(prompt, _model_pair.tokenizer_after, req.system_prompt) + for prompt in req.prompts + ] + + batch_results = [ + LDABatchResult( + prompt=req.prompts[i], + formatted_prompt=formatted_prompts[i], + completion=result.completion, + text=result.text, + tokens_generated=result.tokens_generated, + stopped_early=result.stopped_early, + ) + for i, result in enumerate(results) + ] + + return LDABatchGenerateResponse( + results=batch_results, + alpha=req.alpha, + model_after=MODEL_AFTER_ID, + model_before=MODEL_BEFORE_ID, + batch_size=len(req.prompts), + ) + + if __name__ == "__main__": import uvicorn