Proof of concept LLM chatbot built with Julia: KeemenaLM.jl

KeemenaLM.jl v0.1.0, a Julia proof-of-concept for training and running a small GPT-style chatbot from scratch.

It includes a v9 broad 336M scratch-trained chatbot baseline as a GitHub Release artifact. The model is a research baseline, not a reliable assistant yet, but the Julia training, export, artifact, tokenizer, and REPL path are working.

using Pkg
Pkg.add(url = "https://github.com/mantzaris/KeemenaLM.jl", rev = "v0.1.0")

using KeemenaLM
bundle_dir = download_model("tiny-chatbot-v9-broad-336m")
tokenizer_dir = resolve_tokenizer_bundle("tiny-chatbot-v9-broad-336m")

It also used other Julia packages such as
https://github.com/mantzaris/KeemenaSubwords.jl
and
https://github.com/mantzaris/KeemenaPreprocessing.jl

I would be happy to talk with others interested in Julia for LLMs and work together to improve the training to make a better chatbot, or some other application even.

Looks cool! I am wondering, is it possible to run this on AMD hardware? I currently don’t have access to any NVIDIA GPUs, only MI250X (previous generation, CDNA 2, gfx90a architecture). I’m also fairly new to the topic (Julia: MCP servers and RAG pipelines, Python: basic training using SambaNova models).

I thought it might (because of Flux), but see in the docs:

Device policy

  • v0.1 inference must work on CPU
  • v0.1 inference may also run on NVIDIA GPU
  • v0.1 training is assumed to run on NVIDIA GPU
  • non-NVIDIA GPU support is out of scope for v0.1
  • CPU training is not a product target for v0.1, even if some toy tests may execute there

I’m not sure why Flux wouldn’t handle AMD (since it should?) and thus this package. It seems to use Float32. Maybe it’s only tested on NVIDIA, that he means, not promised elsewhere…

I only have an nvidia gpu but if I get an amd, I would be happy to try to get that to work too. Want to try it out and make a PR?

Do you think this card is any good? I have been using it for several weeks, and AITER is largely non-functional. While configuring libfabric is not overly difficult, it is not entirely straightforward either. On the bright side, these GPUs are plentiful on both sides of the Atlantic. Do you require a large quantity of them for this project or what would be your optimal setup? P.S. And who are you rooting for, guys: Spain or Argentina?

I think it is a good card from the specs but I have not used it or AITER. The nvidia/cuda is just so reliable that it justifies the price hike. I trained on a single 32GB nvidia ADA. The result is a functional chatbot after a week of training, not great though, but it works and demonstrates it can do better. Going forward to the next level would probably need in total 128gb min or 256 to be safe to get a semi decent chatbot rather than a proof of concept with enough storage for all the text examples. I was loosely following the ‘Chinchila’ scaling with the approx 300M param model, and if I was to plan another run it would be 500M (0.5B) on a much bigger dataset.

(ps. I don’t watch football/soccer, so whoever puts on the best performance :slight_smile: )

I just can’t believe it. This is extremely unusual in my region. And you, @Palli, if I may ask, did you see the Spain vs Argentina game? :- )

Any thoughts on the final capacity and hardware requirements, if I may ask? :- )

I am not sure about which capacity you might be referring to but there is plenty of room even at 300M for more training with more data. If you have better infrastructure let me know what you have and we can work on a refactoring that can utilize it better.

Dont feel bad I cant use Enzyme or Reactant because I got an Intel gpu
Nividia is just where the stack is for everything.

I was trying to be useful, and I was wondering about your objective for this model. In other words, what is the scaling goal? It’s different for 300M and 500M, isn’t it? That’s all.

The way I look at it is from the perspective of the objective and the cost basis. Currently, as I informed you, I’m using AMD and it’s pretty neat. However, the latest models don’t work on it. The hardware is too old for them, at least when running via vLLM.

Well, usually. In general, I have a feeling that this field is becoming one where, aside from scale and energy costs, there are almost no barriers to entry. As a result, hardware is starting to be the level playing field. Hence, I think that in general, objectives and hardware are rather important questions here.

You are right that it is going to be different, and depends on the hardware as you say. If I get an AMD gpu I will try to work on its support, but you can work on it too in the mean time. Thinking about it a bit more, with the nvidia costs, maybe putting time just to support amd might be worth it. I can look into getting one around early Sept. There were talks of doing federated learning in julia a few years ago but I don’t recall much activity around that recently.

I think your prudent, low-profile strategy is probably the right approach. I have always preferred Intel and NVIDIA, however, as far as I know, AMD is what our countries have mostly invested in. Even if it’s not at the very top, this hardware is fairly capable and should be able to do the heavy lifting if needed.

At the same time, I share @obsidianjulua’s point of view that it might be useful to have access to the latest NVIDIA stack as well. I believe this is absolutely doable, however, I would prefer not to discuss the details publicly.

I also think it’s worth looking into some of the more exotic solutions available on the market, but focusing on the primary vendors initially would be probably more prudent.

The reason I asked the question about hardware is to avoid a situation where we start working and you suddenly say, oh no, the plan includes things only possible on NVIDIA, or only on AMD, or only on Habana Gaudi, or on any other hardware that might come to mind.

Yesterday, I cloned both the KeemenaLM and NeuroDSL repositories. I will try my best, but please do not expect miracles. I am not a professional coder.

Lastly, if I may ask, how are you planning to shape the project from a legal perspective? I believe this is a crucial question at this stage.

I think I got it running and it seems that currently it is behaving as follows:

After another day with your model, I’d like to confirm my preliminary expression of interest.

Great to see it working, and it ran on your AMD gpu? As you can see, it needs more training, a bigger window, etc. There is work to do!

That is great to hear that you are interested in collaborating! Let’s talk more about making v0.2 better :slight_smile:

So far, my tests have only used CPUs. I haven’t made move_model_to_device fully support AMD yet. Certainly, this is crucial to me because I’m currently on AMD. Also, as I mentioned before, I believe, in general hardware is important. For example, I don’t think it’s possible to do all the quantizations you and @Pali were discussing in the sister thread on the MI250X. It’s simply not feasible. I’d like to emphasize that my questions about the platform were mostly about future developments.

As for the current state of the project, based on my understanding of the code, as for now, the short-term plan would be: 1. Expand behavior eval, 2. Build a direct-answer corpus, 3. Run the current version of the model against the new eval and 4. Re-evaluate and compare. For points 1 and 2, I’m currently considering the following datasets: TruthfulQA, SimpleQA, Natural Questions, TriviaQA, CommonsenseQA, PIQA, ARC-Easy, MMLU, BoolQ, and HellaSwag. However, I haven’t made up my mind about this yet.

You can’t say I’m not. I’ve even tried to draw attention to this project from @csvance, @Khemais_Abdallah and some other persons knowledgeable about ML/AI I’ve met along my Julia way. Also, writing all these posts takes time. However, if I may ask again, how are you planning to shape the project from a legal perspective?

Then maybe the first and most practical goal is to allow AMD users to run on their hardware?

I am not able with my hardware to expand the dataset, as is on a single gpu it takes a week to train. Using AMD might allow my/our budget to get more power into the training. The amd gpus seem so much cheaper, for that price my hardware setup could be more powerful by far.

I asked GPT-5 mini to add support for Metal.jl. After some back and forth, I can run with --device=metal but I cannot tell any difference in speed.

I created a PR with my changes. All tests pass but a broken one, but I am not profficient enough to confirm that all changes are meaningful. Perhaps with a bit of coordinated effort, between us, with one stone we can get all the birds that Flux.jl supports.