gsoc_proposal_content.md

GSoC 2026 Proposal: LLM Support for 7B Models (OLMo) in DeepChem

CONTENT REFERENCE — REWRITE IN YOUR OWN WORDS BEFORE SUBMITTING

---

1. INTRODUCTION

What this project is about

DeepChem's HuggingFaceModel wrapper currently supports encoder-only models (ChemBERTa, MoLFormer) through masked language modeling. There is no support for decoder-only causal language models. This project adds OLMo-2 (Allen AI's open language model) to DeepChem, enabling:

• Continued pretraining on molecular data (SMILES)
• Fine-tuning for classification and regression
• Autoregressive molecular generation

Why it matters

Encoder models (ChemBERTa) can classify and predict properties, but they CANNOT generate molecules. A causal LM like OLMo opens up:

• De novo molecular generation (drug discovery)
• Text-molecule bridging (OLMo understands English AND can learn SMILES)
• In-context few-shot learning without fine-tuning
• Transfer learning from scientific literature

Why OLMo specifically

• Fully open (weights, data, training code) — unlike LLaMA/GPT
• OLMo-2 is natively supported in HuggingFace transformers (no custom code)
• 1B and 7B variants available for different compute budgets
• Trained on Dolma corpus which includes scientific papers

What I've already done (reference your PR and experiments)

• Found and fixed a transformers 5.x compatibility bug in Chemberta (PR #4913)
• Filed issue #4912 documenting broader transformers 5.x compat gap
• Built a working OLMo wrapper prototype (locally) with:
  • Olmo2ForSequenceClassification (doesn't exist in transformers)
  • Causal LM pretraining on SMILES
  • All 8 unit tests passing
• Ran experiments on real MoleculeNet data:
  • BBBP classification: ROC-AUC 0.67 (random init, tiny model)
  • ESOL regression: R^2 = 0.37
  • SMILES generation: 0% validity (expected — proves pretraining is essential)

---

2. RELEVANT EXPERIENCE & INTEREST

Technical background

• Parameter Golf (OpenAI competition, March 2026): Trained language models from scratch under 16MB constraint. Custom SentencePiece tokenizers, GPTQ-lite quantization, flash attention, architecture design (11L 512d transformer). This is directly relevant — I understand transformer training at a low level.
• GSPO-DeepSeek-R1-Distill-Qwen-1.5B (15 GitHub stars): Fine-tuning and distillation of large language models.
• wingman-AI (29 GitHub stars): Production AI assistant system.
• Open source contributions: PRs to HuggingFace transformers, Unsloth, Anthropic SDK, OpenAI SDK, Karpathy's nanochat.

Why I want to work on this

[WRITE THIS YOURSELF — what genuinely interests you about molecular ML? Why DeepChem? Be specific and honest. Don't say "I'm passionate about open source" — say what specific thing drew you to this project.]

Links

• GitHub: https://github.com/vivekvar-dl
• PR #4913: https://github.com/deepchem/deepchem/pull/4913
• Issue #4912: https://github.com/deepchem/deepchem/issues/4912

---

3. WORK PLAN

3.1 Design

The implementation has four components:

Component A: Base class changes to HuggingFaceModel

• Add causal_lm task support (DataCollatorForLanguageModeling with mlm=False)
• Add AutoModelForCausalLM branch in load_from_pretrained()
• Add generate() method for autoregressive text generation
• Add causal LM batch preparation in _prepare_batch()

Note: PR #4907 by another contributor adds a similar generate() method. My work is complementary — I'm adding a full model wrapper, not just generation plumbing.

Component B: Olmo2ForSequenceClassification This class DOES NOT EXIST in HuggingFace transformers. OLMo only has OlmoForCausalLM — no classification head. I built one:

• Extends Olmo2PreTrainedModel
• Uses last-token pooling (last non-padded token's hidden state)
• Linear projection head for classification/regression
• Supports single-label, multi-label, and regression via problem_type config
• Computes CrossEntropyLoss / BCEWithLogitsLoss / MSELoss based on task

This follows the same pattern as LlamaForSequenceClassification.

Component C: OLMo wrapper class

OLMo(HuggingFaceModel)
  __init__(task, model_name, n_tasks, config)
    - task: causal_lm | regression | classification | mtr
    - Loads tokenizer from HuggingFace Hub
    - Sets pad_token = eos_token (decoder models don't have pad by default)
    - Syncs vocab_size between config and tokenizer
    - Creates appropriate model class based on task

  _prepare_batch(batch)
    - causal_lm: labels = input_ids (model shifts internally)
    - regression/classification: labels from dataset, proper dtype casting
    - Multi-task classification: float labels for BCEWithLogitsLoss

Component D: Tokenization strategy Phase 1 (GSoC): Use OLMo's pretrained tokenizer as-is on SMILES.

• OLMo's 100K BPE vocab actually tokenizes SMILES more efficiently than ChemBERTa's 600-token vocab (0.9x token ratio in my analysis)
• BUT it fragments chemical semantics: [C@@H] -> [C, @@, H, ]
• ChemBERTa learns chemistry-aware merges: (=O), ccccc, COc

Phase 2 (stretch): Extend tokenizer with SMILES-specific tokens.

• Add special tokens for stereochemistry: [C@@H], [C@H], [nH]
• Add aromatic ring tokens: c1ccccc1
• Retrain BPE on mixed English + SMILES corpus

3.2 Pseudocode

Olmo2ForSequenceClassification.forward():

  hidden_states = self.model(input_ids, attention_mask)
  # Pool: use last non-padded token
  seq_lengths = (input_ids != pad_token_id).sum(-1) - 1
  pooled = hidden_states[batch_range, seq_lengths]
  logits = self.score(pooled)  # Linear(hidden_size, num_labels)
  if labels:
    if regression: loss = MSELoss(logits, labels)
    if single_class: loss = CrossEntropy(logits, labels)
    if multi_label: loss = BCEWithLogits(logits, labels)
  return {loss, logits}

OLMo._prepare_batch() for causal_lm:

  tokens = tokenizer(smiles_list, padding=True)
  input_ids = tokens.input_ids.to(device)
  labels = input_ids.clone()  # next-token prediction
  return {input_ids, attention_mask, labels}

HuggingFaceModel.generate():

  tokens = tokenizer(inputs, padding=True)
  output_ids = model.generate(**tokens, max_new_tokens=N, **kwargs)
  return tokenizer.batch_decode(output_ids)

3.3 Testing Plan

8 unit tests (all passing in my prototype):

Test	What it validates
test_olmo_causal_lm_pretraining	Causal LM trains, loss > 0
test_olmo_regression_finetuning	Regression trains, predictions match shape, MAE computable
test_olmo_classification	Classification on binary labels, loss > 0
test_olmo_multitask_regression	MTR with 2 tasks, predictions shape matches
test_olmo_save_and_restore	Checkpoint save/load, weights match exactly
test_olmo_load_from_pretrained	Pretrain causal LM -> load into regression model
test_olmo_generate	Single and batch generation returns strings
test_olmo_invalid_task	ValueError on bad task name

All tests use a tiny config (64 hidden, 2 layers, 2 heads) — no model download needed, runs in ~27 seconds on CPU.

Integration tests (to add during GSoC):

• MoleculeNet benchmarks: BBBP, ESOL, FreeSolv, Lipophilicity
• SMILES generation validity (RDKit validation)
• Continued pretraining convergence on ZINC/PubChem subsets

3.4 Sources of Risk

Risk	Likelihood	Mitigation
OLMo-7B requires ~14GB VRAM for inference	High	Use OLMo-1B for CI/demos. Test with tiny configs. Document GPU requirements.
SMILES generation validity low without extensive pretraining	High	This IS the core problem. Budget 3 weeks for pretraining experiments. Use ZINC-250K as training corpus. Target >50% validity.
Olmo2ForSequenceClassification not upstream	Medium	Our implementation follows HF patterns exactly. If HF adds it later, we swap to theirs.
Tokenizer fragments chemical semantics	Medium	Phase 1: works as-is (my experiments show learning happens). Phase 2: extend vocabulary.
transformers version compatibility	Low	Already found and fixed one issue (PR #4913). Use top-level imports throughout.

3.5 Milestones & Timeline

Assuming Medium size (175 hours, ~12 weeks):

Milestone 1: Core wrapper (Weeks 1-3)

• PR: Base class changes to HuggingFaceModel (causal_lm task, generate())
  • Coordinate with PR #4907 to avoid duplication
• PR: Olmo2ForSequenceClassification
• PR: OLMo wrapper class with all task modes
• PR: Unit tests (8 tests)
• Deliverable: from deepchem.models import OLMo works for all tasks

Milestone 2: Continued pretraining (Weeks 4-6)

• PR: Pretraining pipeline on molecular data (ZINC-250K)
• PR: Data loading utilities for SMILES corpora
• PR: Pretraining tutorial notebook
• Deliverable: Pretrained OLMo checkpoint on molecular data

Milestone 3: Fine-tuning & benchmarks (Weeks 7-9)

• PR: Classification tutorial (BBBP, Tox21)
• PR: Regression tutorial (ESOL, FreeSolv, Lipophilicity)
• PR: Benchmark results table vs ChemBERTa
• Deliverable: Published benchmark comparing OLMo vs ChemBERTa on MoleculeNet

Milestone 4: Generation & polish (Weeks 10-12)

• PR: SMILES generation tutorial with RDKit validity checking
• PR: Documentation (numpydoc, API reference, user guide)
• PR: Tokenizer extension experiments (stretch goal)
• Deliverable: Complete documentation and tutorials

Each milestone = 1 evaluation checkpoint. PRs are <50 lines where possible, following DeepChem's contribution guidelines.

3.6 Pull Request Plan

I will follow DeepChem's guidelines: small PRs (<50 lines for initial ones), with tests and numpydoc documentation. Expected ~8-12 PRs total:

1. HuggingFaceModel causal_lm support (~40 lines)
2. generate() method (~50 lines)
3. Olmo2ForSequenceClassification (~100 lines — larger, will discuss with mentor)
4. OLMo wrapper class (~80 lines)
5. Unit tests (~180 lines)
6. Pretraining pipeline
7. Data utilities
8. Tutorial notebooks (3-4 notebooks)
9. Documentation updates
10. Benchmark scripts

---

4. COMMUNITY ENGAGEMENT

• Already contributing: PR #4913 (bug fix), Issue #4912 (compat report)
• Will attend office hours MWF 9am PST
• Will join Discord for async discussion
• Will write weekly progress updates
• Happy to review other contributors' HuggingFace-related PRs

---

5. RESOURCES REQUIRED

• GPU: I have access to 1x H100 NVL 96GB (Azure) for development
• For CI: tiny model configs, no GPU needed
• For pretraining experiments: my H100 is sufficient for OLMo-1B
• OLMo-7B experiments: may need multi-GPU setup (discuss with mentor)

---

6. BIBLIOGRAPHY

1. Groeneveld et al. (2024). "OLMo: Accelerating the Science of Language Models." arXiv:2402.00838
2. Chithrananda et al. (2020). "ChemBERTa: Large-Scale Self-Supervised Pretraining for Molecular Property Prediction." arXiv:2010.09885
3. Ross et al. (2022). "Large-Scale Chemical Language Representations Capture Molecular Structure and Properties." Nature Machine Intelligence.
4. Weininger (1988). "SMILES, a chemical language and information system." J. Chem. Inf. Comput. Sci.
5. Wu et al. (2018). "MoleculeNet: A Benchmark for Molecular Machine Learning." Chemical Science.

---

KEY NUMBERS FROM YOUR EXPERIMENTS (reference these in proposal)

• Tokenization: OLMo uses 0.9x tokens vs ChemBERTa on drug molecules
• BBBP classification: ROC-AUC 0.67 (random init, 12.9M param model, 200 samples, 3 epochs)
• ESOL regression: R^2 = 0.37, MAE = 1.27 (same conditions)
• SMILES generation: 0% validity from random init (proves pretraining is the core challenge)
• Test suite: 8/8 tests pass in 27 seconds on CPU
• Stereochemistry fragmentation: [C@@H] splits into 4 tokens in OLMo vs 7 in ChemBERTa