first commit

examples/scheduler/run_greedy.py

@@ -0,0 +1,153 @@
+#!/usr/bin/env python
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Pre-train or fine-tune a causal language model using GreedyLR or Cosine scheduler.
+
+Example usage:
+
+    # Pre-train with GreedyLR (default):
+    python run_greedy.py
+
+    # Pre-train with cosine scheduler for comparison:
+    python run_greedy.py --lr_scheduler_type cosine
+
+    # Use a different model:
+    python run_greedy.py --model_name_or_path Qwen/Qwen3-0.6B
+
+    # Fine-tune from a pretrained checkpoint:
+    python run_greedy.py --model_name_or_path meta-llama/Llama-3.2-1B --finetune
+"""
+
+import argparse
+import logging
+
+from datasets import load_dataset
+
+from transformers import (
+    AutoConfig,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    DataCollatorForLanguageModeling,
+    Trainer,
+    TrainingArguments,
+    set_seed,
+)
+
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Pre-train/fine-tune a causal LM with GreedyLR or Cosine scheduler")
+    parser.add_argument(
+        "--model_name_or_path",
+        type=str,
+        default="meta-llama/Llama-3.2-1B",
+        help="Model identifier from huggingface.co/models or path to a local checkpoint",
+    )
+    parser.add_argument(
+        "--dataset_name",
+        type=str,
+        default="wikitext",
+        help="The name of the dataset to use (via datasets library)",
+    )
+    parser.add_argument(
+        "--dataset_config_name",
+        type=str,
+        default="wikitext-2-raw-v1",
+        help="The configuration name of the dataset",
+    )
+    parser.add_argument("--lr_scheduler_type", type=str, default="greedy", choices=["greedy", "cosine"])
+    parser.add_argument("--learning_rate", type=float, default=2e-4)
+    parser.add_argument("--max_steps", type=int, default=2000)
+    parser.add_argument("--per_device_train_batch_size", type=int, default=2)
+    parser.add_argument("--gradient_accumulation_steps", type=int, default=8)
+    parser.add_argument("--output_dir", type=str, default="./output")
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument(
+        "--finetune",
+        action="store_true",
+        help="Fine-tune from pretrained weights instead of training from scratch",
+    )
+    parser.add_argument("--block_size", type=int, default=512, help="Context length for tokenization")
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    set_seed(args.seed)
+
+    tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    if args.finetune:
+        model = AutoModelForCausalLM.from_pretrained(args.model_name_or_path)
+    else:
+        config = AutoConfig.from_pretrained(args.model_name_or_path)
+        model = AutoModelForCausalLM.from_config(config)
+
+    param_count = sum(p.numel() for p in model.parameters())
+    logger.info(f"Model: {args.model_name_or_path} ({param_count / 1e6:.1f}M parameters)")
+
+    raw_datasets = load_dataset(args.dataset_name, args.dataset_config_name)
+
+    def tokenize_function(examples):
+        return tokenizer(examples["text"], truncation=True, max_length=args.block_size)
+
+    tokenized_datasets = raw_datasets.map(
+        tokenize_function, batched=True, remove_columns=raw_datasets["train"].column_names
+    )
+
+    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
+
+    scheduler_kwargs = {}
+    if args.lr_scheduler_type == "greedy":
+        scheduler_kwargs = {"patience": 10, "factor": 0.99, "min_lr": 1e-5}
+
+    training_args = TrainingArguments(
+        output_dir=args.output_dir,
+        per_device_train_batch_size=args.per_device_train_batch_size,
+        gradient_accumulation_steps=args.gradient_accumulation_steps,
+        learning_rate=args.learning_rate,
+        lr_scheduler_type=args.lr_scheduler_type,
+        lr_scheduler_kwargs=scheduler_kwargs,
+        max_steps=args.max_steps,
+        warmup_steps=0 if args.lr_scheduler_type == "greedy" else 100,
+        eval_strategy="steps",
+        eval_steps=200,
+        save_steps=500,
+        logging_steps=10,
+        bf16=True,
+        report_to="tensorboard",
+        seed=args.seed,
+    )
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=tokenized_datasets["train"],
+        eval_dataset=tokenized_datasets.get("validation", tokenized_datasets.get("test")),
+        data_collator=data_collator,
+    )
+
+    logger.info(f"Starting training with {args.lr_scheduler_type} scheduler")
+    trainer.train()
+    trainer.save_model()
+
+
+if __name__ == "__main__":
+    main()