Firstly, thank you so much for looking at this post. I could really use some help.

I have followed this guide as closely as possible: https://github.com/kingoflolz/mesh-transformer-jax

I'm trying to fine-tune GPT-J with a small dataset of ~500 lines:

You are important to me. <|endoftext|>
I love spending time with you. <|endoftext|>
You make me smile. <|endoftext|>
feel so lucky to be your friend. <|endoftext|>
You can always talk to me, even if it’s about something that makes you nervous or scared or sad. <|endoftext|>

Using the create_finetune_tfrecords.py script (from the repo mentioned above) outputs a file with 2 in it. I understand that means my data has 2 sequences.

I could really use some advice with the .json config file. What constants do you recommend for this small dataset? The best I came up with trying to follow the guide:

  "layers": 28,
  "d_model": 4096,
  "n_heads": 16,
  "n_vocab": 50400,
  "norm": "layernorm",
  "pe": "rotary",
  "pe_rotary_dims": 64,

  "seq": 2048,
  "cores_per_replica": 8,
  "per_replica_batch": 1,
  "gradient_accumulation_steps": 2,

  "warmup_steps": 1,
  "anneal_steps": 9,
  "lr": 1.2e-4,
  "end_lr": 1.2e-5,
  "weight_decay": 0.1,
  "total_steps": 10,

  "tpu_size": 8,

  "bucket": "chat-app-tpu-bucket-europe",
  "model_dir": "finetune_dir",

  "train_set": "james_bond_1.train.index",
  "val_set": {},

  "eval_harness_tasks": [

  "val_batches": 2,
  "val_every": 400000,
  "ckpt_every": 1,
  "keep_every": 1,

  "name": "GPT3_6B_pile_rotary",
  "wandb_project": "mesh-transformer-jax",
  "comment": ""

Problem is... When I test the fine-tuned model, I get responses that make no sense: Screenshot

Very much looking forward to hearing from you! :)

  • 1
    $\begingroup$ Have you looked at the notes on fine-tuning? Based on what is mentioned it seems that the current learning rate you are using is quite high (1.2e-4 vs a range of 1e-5 and 5e-5). $\endgroup$
    – Oxbowerce
    Nov 17, 2021 at 13:14
  • $\begingroup$ Hmm. I just took that learning rate from an example that was in the repo. I'll try the learning rates you suggested and let you know :) Does everything else in my JSON file look okay though? $\endgroup$ Nov 17, 2021 at 13:48
  • $\begingroup$ I don't have too much knowledge regarding this specific model (type), so I would just compare your settings to what is described on the github page. $\endgroup$
    – Oxbowerce
    Nov 17, 2021 at 14:16

1 Answer 1


If your dataset is small you can train shorter sequence length that will improve the training speed. You can keep the data in multiple lines instead of adding them to the max-token-length.

At the moment it seems that the your data is put in two or more lines of 2048 tokens each. It can be a problem because if the number of epoch = 1, the model will only update twice with the learning rate which means it does not really learn much. If you use batch size=8 then the model simply does not have enough data to batch because the dataset only have 2 or more lines.

Your generated text looks strange can also because of your generation strategies. (refer to this: https://huggingface.co/blog/how-to-generate)

I personally prefer the sampling technique for generation.

Another possible problem is that you use the wrong tokeniser for training or inference which will mess up the model output.


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.