Function parse_args

train_systematicity.py:14–55

()

def parse_args():
    parser = argparse.ArgumentParser(description="Train a RecurrentDepthTransformer model.")

    # Data paths
    parser.add_argument('--data_dir', type=str, default='data/composition.2000.200.7.2',
                        help='Path to dataset directory')
    parser.add_argument('--seed', type=int, default=42)
    parser.add_argument('--checkpoint_dir', type=str, default='checkpoints/systematicity/r4',
                        help='Path to save model checkpoints')
    parser.add_argument('--log_file', type=str, default='results/systematicity/r4.txt',
                        help='Path to save training logs')

    # Training hyperparameters
    parser.add_argument('--num_epochs', type=int, default=150001, help='Number of training epochs')
    parser.add_argument('--batch_size', type=int, default=128, help='Batch size for training and evaluation')
    parser.add_argument('--lr', type=float, default=1e-4, help='Learning rate')
    parser.add_argument('--weight_decay', type=float, default=0.1, help='Weight decay for optimizer')
    parser.add_argument('--warmup_steps', type=int, default=2000,
                        help='Number of warmup steps for learning rate scheduler')

    # Model hyperparameters
    parser.add_argument('--d_model', type=int, default=768, help='Dimension of model embeddings')
    parser.add_argument('--num_recurrent_layers', type=int, default=4, help='Number of recurrent layers')
    parser.add_argument('--num_heads', type=int, default=12, help='Number of attention heads')
    parser.add_argument('--recurrence', type=int, default=4, help='Number of recurrent iterations')

    parser.add_argument('--use_compile', action='store_true',
                        help='Enable torch.compile for the model')
    parser.add_argument('--compile_mode', type=str,
                        choices=['default', 'reduce-overhead', 'max-autotune', 'max-autotune-no-cudagraphs'],
                        default='max-autotune',
                        help='torch.compile mode')

    # Precision settings
    parser.add_argument('--precision', type=str, choices=['fp16', 'bf16'], default='bf16',
                        help='Enable mixed precision training (fp16 or bf16)')

    # Device configuration
    parser.add_argument('--device', type=str, default='cuda' if torch.cuda.is_available() else 'cpu',
                        help='Device to use for training')

    return parser.parse_args()


def train_model(model, dataloader, valid_dataloader, test_dataloader, args):