train/train_deep.py

# import sys
# import torch
# sys.path.append("..")

# import os
# from datasets import load_dataset
# from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments, DataCollatorForLanguageModeling
# from utils_llama import PERTURBATIONS, BABYLM_SPLITS, BABYLM_DATA_PATH, \
#     GENRES, MARKER_TOKEN_IDS, marker_sg_token, marker_pl_token, marker_rev_token, write_file
# import argparse

# os.environ["TOKENIZERS_PARALLELISM"] = "false"

# # import wandb

# # Setup for Weights & Biases
# # wandb.init(project="kallini", group="babylm-perturbation-experiments", name=run_id)

# if __name__ == "__main__":

#     # === CONFIGURATION SETTINGS ===
#     parser = argparse.ArgumentParser(description="Training configuration.")

#     parser.add_argument('--perturbation', type=str, default='hop_tokens4', help='Type of perturbation to use.')
#     parser.add_argument('--train_set', type=str, default='10M', help='Dataset size for training.')
#     parser.add_argument('--batch_size', type=int, default=4, help='Batch size for training.')
#     parser.add_argument('--epoch', type=int, default=20, help='train epoch')
#     parser.add_argument('--seed', type=int, default=0, help='Random seed.')

#     args = parser.parse_args()

#     # no_pos_encodings_underscore = ""  # Ex: "_nopos" if needed
#     ckpt_path = "./checkpoints"
#     # effective_bsz = 512

#     model_name = "meta-llama/Llama-3.2-3B"

#     model_save_name = "Llama-3.2-3B"
#     # === FILE PATHS BASED ON CONFIGURATION ===
#     run_id = f"babylm_{args.perturbation}_{args.train_set}_seed{args.seed}"
#     cache_dir = os.path.join(ckpt_path, f"{model_save_name}", run_id, "artifacts")
#     run_dir = os.path.join(ckpt_path, f"{model_save_name}", run_id, "runs")
#     os.makedirs(cache_dir, exist_ok=True)
#     os.makedirs(run_dir, exist_ok=True)

#     # === DATASET LOADING ===
#     dataset_name = f"babylm_{args.perturbation}_{args.train_set}_seed{args.seed}"
#     # dataset = load_dataset('babylm_dataset_test.py', name=dataset_name, trust_remote_code=True)
#     dataset = load_dataset('babylm_dataset_test.py', name=dataset_name, trust_remote_code=True)
#     train_dataset = dataset['train']
#     val_dataset = dataset['validation']
#     print(train_dataset)    

#     # === TOKENIZER & MODEL LOADING ===
#     # model_name = f"gpt2{'' if no_pos_encodings_underscore == '' else '-no-pos'}-small-{perturbation}-{paren_model}"
 
#     # tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=cache_dir)
#     tokenizer = PERTURBATIONS[args.perturbation]['llama_tokenizer']
#     model = AutoModelForCausalLM.from_pretrained(model_name, 
#                                                 # device_map="auto",  # deepspeed needs to delete this setting
#                                                 cache_dir=cache_dir)

#     # print("model:", model)
#     # === TOKENIZATION ===
#     def tokenize_function(examples):
#         return tokenizer(examples['text'], padding="max_length", truncation=True, max_length=1024)
#     tokenized_train = train_dataset.map(tokenize_function, batched=True, remove_columns=["text"])
#     tokenized_valid = val_dataset.map(tokenize_function, batched=True, remove_columns=["text"])

#     # === DATA COLLATOR ===
#     data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)

#     # === TRAINING ARGUMENTS ===
#     training_args = TrainingArguments(
#         output_dir=run_dir,
#         evaluation_strategy="steps",
#         eval_steps=10,
#         per_device_train_batch_size=args.batch_size,  # set "auto" in deepspeed config, adjust it in trainer
#         logging_dir='./logs',
#         logging_steps=10,
#         save_steps=200000000,
#         learning_rate=5e-5, # align with deepspeed
#         num_train_epochs=args.epoch,
#         seed=args.seed,
#         gradient_accumulation_steps=4, # # set "auto" in deepspeed config, adjust it in trainer
#         fp16 = True, # align with deepspeed
#         report_to="none",
#         deepspeed="deepspeed_config/train_dp_config.json"
#     )

#         # === TRAINER ===
#     trainer = Trainer(
#         model=model,
#         args=training_args,
#         train_dataset=tokenized_train,
#         eval_dataset=tokenized_valid,
#         tokenizer=tokenizer,
#         data_collator=data_collator
#     )
    
#     # === TRAIN MODEL ===
#     trainer.train()
#     # End logging
#     # wandb.finish()

import sys
import torch
sys.path.append("..")

import os
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments, DataCollatorForLanguageModeling
from utils_llama import PERTURBATIONS, BABYLM_SPLITS, BABYLM_DATA_PATH, \
    GENRES, MARKER_TOKEN_IDS, marker_sg_token, marker_pl_token, marker_rev_token, write_file
import argparse

os.environ["TOKENIZERS_PARALLELISM"] = "false"


# Setup for Weights & Biases
# wandb.init(project="kallini", group="babylm-perturbation-experiments", name=run_id)

if __name__ == "__main__":

    # === CONFIGURATION SETTINGS ===
    parser = argparse.ArgumentParser(description="Training configuration.")

    parser.add_argument('--perturbation', type=str, default='hop_tokens4', help='Type of perturbation to use.')
    parser.add_argument('--train_set', type=str, default='10M', help='Dataset size for training.')
    parser.add_argument('--batch_size', type=int, default=4, help='Batch size for training.')
    parser.add_argument('--epoch', type=int, default=20, help='train epoch')
    parser.add_argument('--seed', type=int, default=0, help='Random seed.')

    args = parser.parse_args()

    # no_pos_encodings_underscore = ""  # Ex: "_nopos" if needed
    ckpt_path = "./checkpoints"
    # effective_bsz = 512

    model_name = "meta-llama/Llama-3.2-3B"

    model_save_name = "Llama-3.2-3B"
    # === FILE PATHS BASED ON CONFIGURATION ===
    run_id = f"babylm_{args.perturbation}_{args.train_set}_seed{args.seed}"
    cache_dir = os.path.join(ckpt_path, f"{model_save_name}", run_id, "artifacts")
    run_dir = os.path.join(ckpt_path, f"{model_save_name}", run_id, "runs")
    os.makedirs(cache_dir, exist_ok=True)
    os.makedirs(run_dir, exist_ok=True)

    # === DATASET LOADING ===
    dataset_name = f"babylm_{args.perturbation}_{args.train_set}_seed{args.seed}"
    # dataset = load_dataset('babylm_dataset_llama.py', name=dataset_name)
    dataset = load_dataset('babylm_dataset_llama.py', name=dataset_name, trust_remote_code=True)
    train_dataset = dataset['train']

    # === TOKENIZER & MODEL LOADING ===
    # model_name = f"gpt2{'' if no_pos_encodings_underscore == '' else '-no-pos'}-small-{perturbation}-{paren_model}"
 
    # tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=cache_dir)
    tokenizer = PERTURBATIONS[args.perturbation]['llama_tokenizer']
    model = AutoModelForCausalLM.from_pretrained(model_name, 
                                                # device_map="auto",  # deepspeed needs to delete this setting
                                                cache_dir=cache_dir)

    # print("model:", model)
    # === TOKENIZATION ===
    def tokenize_function(examples):
        return tokenizer(examples['text'], padding="max_length", truncation=True, max_length=1024)
    tokenized_train = train_dataset.map(tokenize_function, batched=True, remove_columns=["text"])

    # === DATA COLLATOR ===
    data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)

    # === TRAINING ARGUMENTS ===
    training_args = TrainingArguments(
        output_dir=run_dir,
        evaluation_strategy="no",
        per_device_train_batch_size=args.batch_size,  # set "auto" in deepspeed config, adjust it in trainer
        logging_dir='./logs',
        logging_steps=150,
        save_steps=10000,
        learning_rate=5e-5, # align with deepspeed
        num_train_epochs=args.epoch,
        seed=args.seed,
        gradient_accumulation_steps=4, # # set "auto" in deepspeed config, adjust it in trainer
        fp16 = True, # align with deepspeed
        report_to="none",
        deepspeed="deepspeed_config/train_dp_config.json"
    )

        # === TRAINER ===
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=tokenized_train,
        tokenizer=tokenizer,
        data_collator=data_collator
    )
    
    # === TRAIN MODEL ===
    trainer.train()
    # End logging
    # wandb.finish()