Modelling the molecular gas that is detected through CO observations of high-z galaxies constitutes a major challenge for ab initio simulations of galaxy formation. I will present recent numerical work aimed at studying the formation and evolution of the simplest and most abundant molecule, H2. Our model fully solves the out-of-equilibrium rate equations and accounts for the unresolved structure of molecular clouds. We apply our model to two types of cosmological simulations: a) the formation of a Milky Way-sized galaxy at z=2 and b) a small cosmological box in order to obtain some statistical results. The results are compared to those obtained from two different approximations commonly used in the literature and for numerical convergence. Our results indicate that independently of the model, robust results (H2 masses) can only be obtained for galaxies that are suffiiciently metal enriched in which H2 formation is fast. However, their morphology differ from model to model. Furthermore, the cosmological H2 mass function derived from the non-equilibrium model agrees well with recents observations that only sample the high-mass end. Extensions of our model towards including other molecules, such as CO, and species, in particular C and its derivatives, will also be discussed.