first commit

examples/scheduler/README.md

@@ -0,0 +1,240 @@
+# GreedyLR: Adaptive Learning Rate Scheduler
+
+GreedyLR monitors training metrics and adaptively adjusts the learning rate -- increasing when improving, decreasing when plateauing. It works for both pre-training and fine-tuning.
+
+## Paper
+
+> **GreedyLR: A Novel Adaptive Learning Rate Scheduler**
+>
+> Despite significant advances in optimizers for training, most research works use common scheduler choices like Cosine or exponential decay. In this paper, we study GreedyLR, a novel scheduler that adaptively adjusts the learning rate during training based on the current loss. To validate the effectiveness of our proposed scheduler, we conduct experiments on several NLP, CV, and LLM tasks with up to 7B parameters, including both fine-tuning and pretraining experiments. The results show that our approach outperforms several state-of-the-art schedulers in terms of accuracy, speed, and convergence.
+
+arXiv: https://arxiv.org/abs/2512.14527
+
+## How It Works
+
+```
++-------------------------------------------------------------+
+|                    GreedyLR Decision Flow                    |
+|                                                              |
+|   Metrics Improving?  --Yes-->  Increase LR (/ factor)       |
+|         |                              |                     |
+|         No                        Enter Warmup               |
+|         |                              |                     |
+|         v                              v                     |
+|   Metrics Plateau?   --Yes-->  Decrease LR (* factor)        |
+|         |                              |                     |
+|         No                        Enter Cooldown             |
+|         |                              |                     |
+|         v                              v                     |
+|   Continue Training              Continue Training           |
++-------------------------------------------------------------+
+```
+
+### Key Parameters
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `patience` | 10 | Steps to wait before adjusting LR |
+| `factor` | 0.95 | Multiplicative factor for LR adjustment |
+| `min_lr` | 1e-3 | Minimum learning rate bound |
+| `smooth` | False | Apply streaming average to metrics |
+
+## Quick Start
+
+### Fine-tuning LLaMA 3.2 1B on simpleCoT
+
+Fine-tuning runs on a single GPU. This example uses [w601sxs/simpleCoT](https://huggingface.co/datasets/w601sxs/simpleCoT), a chain-of-thought reasoning dataset:
+
+```python
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    DataCollatorForLanguageModeling,
+    Trainer,
+    TrainingArguments,
+)
+from datasets import load_dataset
+
+model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B", dtype="bfloat16")
+tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B")
+tokenizer.pad_token = tokenizer.eos_token
+
+raw = load_dataset("w601sxs/simpleCoT", split="train[:5000]")
+eval_raw = load_dataset("w601sxs/simpleCoT", split="train[5000:5500]")
+
+def format_and_tokenize(examples):
+    texts = [
+        f"Question: {s}\nReasoning: {r}\nAnswer: {t}"
+        for s, r, t in zip(examples["source"], examples["rationale"], examples["target"])
+    ]
+    return tokenizer(texts, truncation=True, max_length=512, padding=False)
+
+train_ds = raw.map(format_and_tokenize, batched=True, remove_columns=raw.column_names)
+eval_ds = eval_raw.map(format_and_tokenize, batched=True, remove_columns=eval_raw.column_names)
+
+training_args = TrainingArguments(
+    output_dir="./output",
+    per_device_train_batch_size=2,
+    gradient_accumulation_steps=4,
+    learning_rate=5e-5,
+    lr_scheduler_type="greedy",
+    lr_scheduler_kwargs={"patience": 2, "factor": 0.9, "min_lr": 1e-7},
+    max_steps=2000,
+    eval_strategy="steps",
+    eval_steps=100,
+    bf16=True,
+)
+
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=train_ds,
+    eval_dataset=eval_ds,
+    data_collator=DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False),
+)
+trainer.train()
+```
+
+### Pre-training LLaMA 3.2 1B on RedPajama (multi-GPU)
+
+Pre-training a 1B+ parameter model requires multiple GPUs. This example uses a subset of [RedPajama-Data-V2](https://huggingface.co/datasets/togethercomputer/RedPajama-Data-V2) with DeepSpeed ZeRO-3 across 4 GPUs. The dataset should be pre-tokenized with 2048-token sequences.
+
+Create a DeepSpeed config (`ds_config.json`):
+
+```json
+{
+    "bf16": {"enabled": true},
+    "zero_optimization": {
+        "stage": 3,
+        "offload_optimizer": {"device": "none"},
+        "offload_param": {"device": "none"},
+        "overlap_comm": true,
+        "contiguous_gradients": true,
+        "stage3_gather_16bit_weights_on_model_save": true
+    },
+    "gradient_accumulation_steps": "auto",
+    "gradient_clipping": 1.0,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto"
+}
+```
+
+Training script (`pretrain_greedy.py`):
+
+```python
+import os
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+
+from transformers import (
+    AutoConfig,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    DataCollatorForLanguageModeling,
+    Trainer,
+    TrainingArguments,
+)
+from datasets import load_from_disk
+
+config = AutoConfig.from_pretrained("meta-llama/Llama-3.2-1B")
+model = AutoModelForCausalLM.from_config(config)
+tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B")
+tokenizer.pad_token = tokenizer.eos_token
+
+tokenized_datasets = load_from_disk("./datasets/redpajama")
+
+training_args = TrainingArguments(
+    output_dir="./output",
+    per_device_train_batch_size=1,
+    per_device_eval_batch_size=1,
+    gradient_accumulation_steps=32,
+    learning_rate=2e-4,
+    weight_decay=0.1,
+    max_grad_norm=1.0,
+    lr_scheduler_type="greedy",
+    lr_scheduler_kwargs={"patience": 2, "factor": 0.95, "min_lr": 1e-5},
+    max_steps=2000,
+    eval_strategy="steps",
+    eval_steps=500,
+    bf16=True,
+    gradient_checkpointing=True,
+    dataloader_drop_last=True,
+    deepspeed="ds_config.json",
+)
+
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=tokenized_datasets["train"],
+    eval_dataset=tokenized_datasets["valid"],
+    data_collator=DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False),
+)
+trainer.train()
+```
+
+Launch with `torchrun`:
+
+```bash
+torchrun --nproc_per_node=4 pretrain_greedy.py
+```
+
+### Standalone Usage (without Trainer)
+
+```python
+import torch
+from transformers import GreedyLR
+
+optimizer = torch.optim.AdamW(model.parameters(), lr=2e-4)
+scheduler = GreedyLR(
+    optimizer,
+    mode="min",
+    factor=0.99,
+    patience=10,
+    min_lr=1e-5,
+    smooth=True,
+    window_size=50,
+)
+
+for epoch in range(num_epochs):
+    train_loss = train_one_epoch(model, train_loader, optimizer)
+    val_loss = evaluate(model, val_loader)
+    scheduler.step(val_loss)
+```
+
+## Comparison with Other Schedulers
+
+| Feature | Cosine | ReduceLROnPlateau | GreedyLR |
+|---------|--------|-------------------|----------|
+| Adaptive to metrics | No | Yes (decrease only) | Yes (increase and decrease) |
+| LR increase | No | No | Yes |
+| Warmup after increase | No | No | Yes |
+| Metric smoothing | No | No | Yes (optional) |
+| Auto-reset | No | No | Yes |
+
+## Using with Other Models
+
+GreedyLR is model-agnostic. The `run_greedy.py` script pre-trains or fine-tunes a causal language model on WikiText-2:
+
+```bash
+# Pre-train with GreedyLR (default)
+python examples/scheduler/run_greedy.py
+
+# Compare with cosine scheduler
+python examples/scheduler/run_greedy.py --lr_scheduler_type cosine
+
+# Use a different model
+python examples/scheduler/run_greedy.py --model_name_or_path Qwen/Qwen3-0.6B
+
+# Fine-tune from pretrained weights
+python examples/scheduler/run_greedy.py --finetune
+```
+
+## Citation
+
+```bibtex
+@article{greedylr2025,
+  title={GreedyLR: A Novel Adaptive Learning Rate Scheduler},
+  author={Subramanian, Shreyas and Krishnamoorthy, Bala and Murthy, Pranav},
+  journal={arXiv preprint arXiv:2512.14527},
+  year={2025}
+}
+```