---

name: palm-modeling description: "Operate the PaLM transformer core in palm-rlhf-pytorch: construction, loss/generation checks, LoRA finetuning scopes, flash-attention choices, and reference-only pretraining." disable-model-invocation: true metadata: disco-role: operating license: MIT

PaLM Modeling

Use this sub-skill when a task is about the package's base PaLM transformer rather than reward models or policy-optimization trainer loops.

Route Here For

• Constructing a PaLM model for autoregressive language modeling or non-causal embedding use.
• Running a tiny loss/backward/logits/embedding/generation smoke check with scripts/tiny_palm_smoke.py.
• Understanding PaLM.forward, PaLM.generate, sampling helpers, embedding/logit return modes, save/load, dropout changes, and LoRA finetune scopes.
• Deciding whether to enable flash_attn=True; this package uses PyTorch scaled-dot-product attention, not the external flash-attn package.
• Adapting the repository's PaLM pretraining flow as a reference-only recipe without running the expensive source training script.

Do Not Use This For

• Reward scoring, binned rewards, prompt masks, or implicit process rewards; route those to the reward-modeling sub-skill.
• PPO, GRPO, TPO, FlowRL, RLHFTrainer, or post-training loops; route those to the policy-optimization sub-skill.
• Full-scale enwik8 training or benchmark-quality model recovery. This sub-skill covers mechanics and bounded checks, not long training.

Start Here

1. For API facts and shapes, read references/api-reference.md.
2. For the distilled pretraining recipe and why the original training flow is reference-only, read references/pretraining-workflows.md.
3. For common failures, read references/troubleshooting.md.
4. To verify an installed package quickly, run the bundled helper instead of any source-repo script:

python sub-skills/palm-modeling/scripts/tiny_palm_smoke.py --device auto --check-lora

If running from inside this sub-skill directory, use:

python scripts/tiny_palm_smoke.py --device auto --check-lora

Operating Notes

• Prefer device-neutral examples: create a torch.device, move model and tensors with .to(device), and avoid README-style unconditional .cuda() calls on CPU hosts.
• Keep generation targets explicit: generate(seq_len, prompt=...) treats seq_len as the target total length. With the default return_seq_without_prompt=True, the returned tensor is only the generated suffix of length seq_len - prompt_len when seq_len > prompt_len.
• Use palm.palm_parameters() for base pretraining optimizers and palm.finetune_parameters(scope) only after creating a matching LoRA scope.
• flash_attn=True compares the torch SDPA backend, not a separate flash-attn dependency. Treat it as an optional acceleration choice.
• The base model can return logits, embeddings, or both; use the tiny smoke helper to confirm the shapes before debugging a larger workflow.
• The source training recipe is reference-only because it is long-running and depends on a missing extra in the repository metadata.

Common Prompt Patterns

• "Build a small causal LM smoke test" → use the bundled tiny smoke helper.
• "Check generation length or logits shape" → use return_logits_with_embedding=True or return_seq_without_prompt=False.
• "Add and merge a LoRA scope" → create a fresh unique scope, train only finetune_parameters(scope), then merge intentionally.
• "Should I set flash_attn=True?" → yes only when you are intentionally comparing the SDPA path.
• "Can I run the pretraining script unchanged?" → not as a default runtime workflow; treat it as reference-only and note the missing lion_pytorch extra.

Quick Decision Points

• If the task is about next-token loss or embedding shapes, use a tiny PaLM smoke before anything larger.
• If the task is about device portability, use .to(device) in examples and avoid unconditional .cuda() calls.
• If the task is about LoRA parameters, keep base and adapter optimizers separate unless you explicitly want a merged model.
• If the task is about pretrained weights, remember the repository does not ship a checkpoint.
• If the task is about long training, point the user to the pretraining reference and explain that it is not a smoke path.

Tiny Verification Checklist

A correct tiny check should:

• instantiate a very small PaLM on the selected device;
• compute a scalar loss and call backward();
• confirm logits shape (batch, seq_len, num_tokens);
• confirm embedding shape (batch, seq_len, dim);
• confirm suffix length seq_len - prompt_len when the prompt is omitted from the return;
• optionally add, use, remove, and merge a LoRA scope.

Related Files

• references/api-reference.md for verified signatures and generation semantics.
• references/pretraining-workflows.md for the distilled training recipe and why it is reference-only.
• references/troubleshooting.md for CPU/device, flash-attn, generation, and LoRA issues.
• scripts/tiny_palm_smoke.py for the safe runnable smoke check.