Thanks for following up!
I’m using my own train.py based on this example. It’s been modified from the example to replicate a setup I’ve been using before moving to SageMaker, specifically:


Use the DistilBertTokenizer from a custom vocab file (instead of AutoTokenizer)


Include DataCollatorForLanguageModeling for fine-tuning a masked language model


Replicate previous DistilBertConfig parameters (instead of AutoModelForMaskedLM)


However, beyond that, it largely follows the format from the example - so the saving strategy should be the same/default, which is presumably save_strategy="steps". Is there a way for me to to attach my train.py instead of just pasting it in plain text on here?
Anyway, the tokenized dataset is around 13 GB in total - processed in a separate step and saved as an .arrow file prior to submitting the training job.
Here are the hyper parameters I’m using:
hyperparameters={'model_name':'distilbert-base-uncased',
                 'epochs': 2,
                 'train_batch_size': 32, # 32 with ml.p4d.24xlarge
                 'eval_batch_size': 32,
                 }
And here is the model config:
config = DistilBertConfig(
    vocab_size = tokenizer.vocab_size, 
    max_position_embeddings=2048, # 512 default
    sinusoidal_pos_embds=False, 
    n_layers=6, 
    n_heads=12, 
    dim=768, 
    hidden_dim=3072, 
    dropout=0.1, 
    attention_dropout=0.1, 
    activation='gelu', 
    initializer_range=0.02, 
    qa_dropout=0.1, 
    seq_classif_dropout=0.2, 
    pad_token_id=0
    )
So based on what you’re saying, I’ll try changing the saving strategy to save_strategy="epoch" to reduce the number of checkpoints. To save just the last model in opt/ml/model, is it just as simple as passing save_total_limit=1 in the training argument?
I do have the CloudWatch logs, however I probably can’t paste it here since it’d be way too long. Is there a way I can send that directly to you or the AWS team you’re in contact with? (And if it’s indeed a problem with the data agent, do you have any idea on how I can resolve the problem?)
Thanks so much for helping!